refactor(tui): anchor turn token usage on idle markers (#49001)

packages/tui/src/routes/session/index.tsx

@@ -86,6 +86,7 @@ import { usePlugin } from "../../plugin/context"
 import {
   cacheReuseDrop,
   createSessionRows,
+  legacyTurns,
   messageBoundaryIDs,
   resolvePart,
   sessionRowID,
@@ -153,6 +154,7 @@ export function Session(props: {
   const session = createMemo(() => data.session.get(route.sessionID))
   const messages = () => data.session.message.list(route.sessionID)
   const messageIndexes = createMemo(() => new Map(messages().map((message, index) => [message.id, index])))
+  const legacy = createMemo(() => legacyTurns(messages()))
   const messagesBeforeRevert = () => {
     const messageID = session()?.revert?.messageID
     if (!messageID) return messages()
@@ -1256,6 +1258,7 @@ export function Session(props: {
         diffWrapMode,
         models,
         messageIndex: (messageID) => messageIndexes().get(messageID),
+        legacyTurns: legacy,
         config,
         mutatePending,
         pendingDelivery: (inboxID) => pendingDeliveries().get(inboxID),
@@ -1937,9 +1940,11 @@ function AssistantFooter(props: { message: SessionMessageAssistant }) {
         ?.name ?? `${props.message.model.providerID}/${props.message.model.id}`,
   )
   const messages = createMemo(() => data.session.message.list(ctx.sessionID))
-  const duration = createMemo(() => turnDuration(props.message, messages(), ctx.messageIndex(props.message.id)))
+  const duration = createMemo(() =>
+    turnDuration(props.message, messages(), ctx.messageIndex(props.message.id), ctx.legacyTurns()),
+  )
   const tokensPerSecond = createMemo(() =>
-    turnTokensPerSecond(props.message, messages(), ctx.messageIndex(props.message.id)),
+    turnTokensPerSecond(props.message, messages(), ctx.messageIndex(props.message.id), ctx.legacyTurns()),
   )
   const interrupted = createMemo(() => props.message.error?.message === "Step interrupted")
   return (

packages/tui/src/routes/session/render-context.tsx

@@ -22,6 +22,8 @@ export const context = createContext<{
   diffWrapMode: () => "word" | "none"
   models: () => ModelInfo[]
   messageIndex: (messageID: string) => number | undefined
+  /** True when the session has no idle markers, so turn footers end at the next prompt. */
+  legacyTurns: () => boolean
   config: ReturnType<typeof useConfig>["data"]
   mutatePending: (action: PendingAction, inboxID: string) => Promise<boolean>
   pendingDelivery: (inboxID: string) => SessionInbox.Delivery | undefined

packages/tui/src/routes/session/rows.ts

@@ -123,7 +123,7 @@ export function createSessionRows(sessionID: Accessor<string>, onSynced?: (sessi
                   input: data.session.input.has(sessionID(), message.id),
                 },
               ]
-            : message.type === "compaction"
+            : message.type === "compaction" || message.type === "idle"
               ? [
                   {
                     id: message.id,
@@ -283,13 +283,34 @@ export function reduceSessionRows(messages: SessionMessageInfo[], inputs = new S
   const usage = turnTokens
     ? { steps: [] as SessionMessageAssistant[], previousTurnCache: undefined as CacheUsage | undefined }
     : undefined
+  // Without any idle marker, history predates markers and a turn ends with its terminal step.
+  // With markers, steers keep the turn open, so usage accumulates until the marker closes it.
+  const legacy = !messages.some((message) => message.type === "idle")
+  const flushTurn = (rows: ProjectionEntry[]) => {
+    if (!usage) return
+    const steps = usage.steps.filter(hasTokenUsage)
+    const last = steps.at(-1)
+    usage.steps.length = 0
+    if (!last) return
+    rows.push({
+      entry: {
+        type: "turn-usage",
+        messageIDs: steps.map((step) => step.id),
+        ...(usage.previousTurnCache === undefined ? {} : { previousCache: usage.previousTurnCache }),
+      },
+    })
+    usage.previousTurnCache = { read: last.tokens.cache.read, model: last.model }
+  }
   const entries = [
     ...messages.filter((message) => !pending.has(message.id)),
     ...pendingCompactions,
     ...messages.filter(isInput),
   ].reduce<ProjectionEntry[]>((rows, message) => {
     if (message.type !== "assistant") {
-      if (message.type === "idle") return rows
+      if (message.type === "idle") {
+        flushTurn(rows)
+        return rows
+      }
       if (message.type === "synthetic" && !message.description?.trim()) return rows
       if (message.type === "compaction" && message.status === "completed" && usage) usage.previousTurnCache = undefined
       rows.push({ entry: { type: "message", messageID: message.id }, closesPrevious: !pending.has(message.id) })
@@ -306,21 +327,7 @@ export function reduceSessionRows(messages: SessionMessageInfo[], inputs = new S
     if (terminal || message.retry) {
       rows.push({ entry: { type: "assistant-footer", messageID: message.id } })
     }
-    if (terminal && usage) {
-      const stepsWithUsage = usage.steps.filter(hasTokenUsage)
-      const last = stepsWithUsage.at(-1)
-      if (last) {
-        rows.push({
-          entry: {
-            type: "turn-usage",
-            messageIDs: stepsWithUsage.map((step) => step.id),
-            ...(usage.previousTurnCache === undefined ? {} : { previousCache: usage.previousTurnCache }),
-          },
-        })
-        usage.previousTurnCache = { read: last.tokens.cache.read, model: last.model }
-      }
-      usage.steps.length = 0
-    }
+    if (terminal && legacy) flushTurn(rows)
     return rows
   }, [])
   return projectEntries(entries)
@@ -340,10 +347,18 @@ export function cacheReuseDrop(previous: CacheUsage | undefined, current: CacheU
   return drop > 0 ? drop : undefined
 }
 
-export function turnDuration(message: SessionMessageAssistant, messages: SessionMessageInfo[], position?: number) {
+// `legacy` marks a session without idle markers, where a turn ends at the next prompt. Reactive
+// callers should pass it from a shared memo: the default scans every message, which subscribes
+// the footer to the whole history.
+export function turnDuration(
+  message: SessionMessageAssistant,
+  messages: SessionMessageInfo[],
+  position?: number,
+  legacy = legacyTurns(messages),
+) {
   if (message.time.completed === undefined) return 0
   const index = position ?? messages.findIndex((item) => item.id === message.id)
-  const input = messages[inputIndex(messages, index === -1 ? messages.length : index)]
+  const input = messages[inputIndex(messages, index === -1 ? messages.length : index, legacy)]
   return Math.max(0, message.time.completed - (input?.time.created ?? message.time.created))
 }
 
@@ -351,10 +366,11 @@ export function turnTokensPerSecond(
   message: SessionMessageAssistant,
   messages: SessionMessageInfo[],
   position?: number,
+  legacy = legacyTurns(messages),
 ) {
   const index = position ?? messages.findIndex((item) => item.id === message.id)
   const end = index === -1 ? messages.length : index + 1
-  const start = inputIndex(messages, end)
+  const start = inputIndex(messages, end, legacy)
   const steps = messages
     .slice(start + 1, end)
     .filter((item): item is SessionMessageAssistant => item.type === "assistant")
@@ -368,13 +384,23 @@ export function turnTokensPerSecond(
   return output / (duration / 1_000)
 }
 
-function inputIndex(messages: SessionMessageInfo[], end: number) {
-  // Reading a sliced prefix subscribes every footer to unrelated historical messages.
+export function legacyTurns(messages: SessionMessageInfo[]) {
+  return !messages.some((message) => message.type === "idle")
+}
+
+function inputIndex(messages: SessionMessageInfo[], end: number, legacy: boolean) {
+  // Reading a sliced prefix subscribes every footer to unrelated historical messages, so walk
+  // back only as far as the turn boundary: the nearest input in legacy sessions, otherwise the
+  // first input after the previous idle marker.
+  let input = -1
   for (let index = end - 1; index >= 0; index--) {
     const message = messages[index]
-    if (message.type === "user" || message.type === "synthetic") return index
+    if (message.type === "idle") return input
+    if (message.type !== "user" && message.type !== "synthetic") continue
+    if (legacy) return index
+    input = index
   }
-  return -1
+  return input
 }
 
 function hasTokenUsage(

packages/tui/test/cli/tui/session-rows.test.ts

@@ -226,6 +226,77 @@ test("resets the cross-turn cache baseline after compaction", () => {
   ])
 })
 
+test("closes turn usage on the idle marker so steered steps share one footer", () => {
+  const step = (id: string, read: number) => ({
+    ...assistant(id, []),
+    finish: "stop" as const,
+    tokens: { input: 1, output: 0, reasoning: 0, cache: { read, write: 0 } },
+  })
+  const idle = (id: string, created: number): SessionMessageInfo => ({
+    type: "idle",
+    id,
+    outcome: "succeeded",
+    time: { created },
+  })
+  const messages: SessionMessageInfo[] = [
+    { type: "user", id: "user-1", text: "First", time: { created: 0 } },
+    step("assistant-1", 1_000),
+    { type: "user", id: "steer", text: "Also this", time: { created: 2 } },
+    step("assistant-2", 2_000),
+    idle("idle-1", 3),
+    { type: "user", id: "user-2", text: "Second", time: { created: 4 } },
+    step("assistant-3", 3_000),
+    { type: "user", id: "steer-2", text: "Wait", time: { created: 5 } },
+    step("assistant-4", 4_000),
+  ]
+
+  expect(reduceSessionRows(messages, new Set(), true)).toEqual([
+    { type: "message", messageID: "user-1" },
+    { type: "assistant-footer", messageID: "assistant-1" },
+    { type: "message", messageID: "steer" },
+    { type: "assistant-footer", messageID: "assistant-2" },
+    { type: "turn-usage", messageIDs: ["assistant-1", "assistant-2"] },
+    { type: "message", messageID: "user-2" },
+    { type: "assistant-footer", messageID: "assistant-3" },
+    { type: "message", messageID: "steer-2" },
+    { type: "assistant-footer", messageID: "assistant-4" },
+  ])
+  expect(
+    reduceSessionRows([...messages, idle("idle-2", 6)], new Set(), true).filter((row) => row.type === "turn-usage"),
+  ).toEqual([
+    { type: "turn-usage", messageIDs: ["assistant-1", "assistant-2"] },
+    {
+      type: "turn-usage",
+      messageIDs: ["assistant-3", "assistant-4"],
+      previousCache: { read: 2_000, model: { id: "model", providerID: "provider" } },
+    },
+  ])
+})
+
+test("measures a marker-era turn from its first prompt across steers", () => {
+  const step = (id: string, created: number, streamed: number, completed: number, output: number) => ({
+    ...assistant(id, []),
+    time: { created, streamed, completed },
+    tokens: { input: 1, output, reasoning: 0, cache: { read: 0, write: 0 } },
+  })
+  const messages: SessionMessageInfo[] = [
+    { type: "user", id: "old-input", text: "Old question", time: { created: 0 } },
+    step("old-step", 1_000, 2_000, 3_000, 5),
+    { type: "idle", id: "idle-1", outcome: "succeeded", time: { created: 3_500 } },
+    { type: "user", id: "input", text: "Question", time: { created: 4_000 } },
+    step("first-step", 5_000, 6_000, 7_000, 10),
+    { type: "user", id: "steer", text: "Also this", time: { created: 7_500 } },
+    step("second-step", 8_000, 9_000, 10_000, 20),
+  ]
+  const final = messages[6]
+  if (final.type !== "assistant") throw new Error("Expected an assistant")
+
+  expect(turnDuration(final, messages)).toBe(6_000)
+  expect(turnTokensPerSecond(final, messages)).toBe(15)
+  expect(turnDuration(final, messages, 6, true)).toBe(2_500)
+  expect(turnTokensPerSecond(final, messages, 6, true)).toBe(20)
+})
+
 test("assigns assistant boundaries to the first rendered row instead of the first text row", () => {
   const messages: SessionMessageInfo[] = [
     { type: "user", id: "user-1", text: "Question", time: { created: 0 } },