refactor(tui): simplify turn usage reduction (#38514)

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

@@ -1036,16 +1036,6 @@ type SessionRowViewProps = {
 }
 
 function SessionRowView(props: SessionRowViewProps) {
-  const config = useConfig()
-  const hidden = () => props.row.type === "turn-usage" && config.data.debug?.turn_tokens !== true
-  return (
-    <Show when={!hidden()}>
-      <SessionRowContent row={props.row} message={props.message} boundaryID={props.boundaryID} />
-    </Show>
-  )
-}
-
-function SessionRowContent(props: SessionRowViewProps) {
   return (
     <box id={props.boundaryID} marginTop={1} flexShrink={0}>
       <Switch>

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

@@ -1,6 +1,7 @@
 import type { SessionMessageAssistant, SessionMessageInfo } from "@opencode-ai/client"
 import { createEffect, on, onCleanup, type Accessor } from "solid-js"
 import { createStore, produce, reconcile } from "solid-js/store"
+import { useConfig } from "../../config"
 import { useData } from "../../context/data"
 import { useClient } from "../../context/client"
 
@@ -32,14 +33,20 @@ export type SessionRow =
 export function createSessionRows(sessionID: Accessor<string>) {
   const data = useData()
   const client = useClient()
+  const config = useConfig()
   const [rows, setRows] = createStore<SessionRow[]>([])
   const revertBoundary = () => data.session.get(sessionID())?.revert?.messageID
+  const turnTokens = () => config.data.debug?.turn_tokens === true
 
   function reduce() {
     const messages = data.session.message.list(sessionID())
     const inputs = new Set(data.session.input.list(sessionID()))
     const boundary = revertBoundary()
-    const rows = reduceSessionRows(boundary ? messages.filter((message) => message.id < boundary) : messages, inputs)
+    const rows = reduceSessionRows(
+      boundary ? messages.filter((message) => message.id < boundary) : messages,
+      inputs,
+      turnTokens(),
+    )
     partitionPending(rows, pendingPermissions())
     const position = rows.findIndex((row) => row.type === "message" && inputs.has(row.messageID))
     rows.splice(
@@ -129,23 +136,7 @@ export function createSessionRows(sessionID: Accessor<string>) {
   )
 
   createEffect(
-    on(
-      () =>
-        data.session.message.list(sessionID()).flatMap((message) =>
-          message.type === "assistant"
-            ? [
-                {
-                  id: message.id,
-                  finish: message.finish,
-                  error: message.error,
-                  retry: message.retry,
-                  tokens: message.tokens,
-                },
-              ]
-            : [],
-        ),
-      () => setRows(reconcile(reduce())),
-    ),
+    on(turnTokens, () => setRows(reconcile(reduce()))),
   )
 
   const appendMessage = (messageID: string) =>
@@ -267,9 +258,12 @@ export function createSessionRows(sessionID: Accessor<string>) {
     data.on("session.step.ended", (event) => {
       if (event.data.sessionID !== sessionID() || ["tool-calls", "unknown"].includes(event.data.finish)) return
       appendFooter(event.data.assistantMessageID)
+      if (turnTokens()) setRows(reconcile(reduce()))
     }),
     data.on("session.step.failed", (event) => {
-      if (event.data.sessionID === sessionID()) appendFooter(event.data.assistantMessageID)
+      if (event.data.sessionID !== sessionID()) return
+      appendFooter(event.data.assistantMessageID)
+      if (turnTokens()) setRows(reconcile(reduce()))
     }),
   ]
   onCleanup(() => subscriptions.forEach((unsubscribe) => unsubscribe()))
@@ -277,14 +271,17 @@ export function createSessionRows(sessionID: Accessor<string>) {
   return rows
 }
 
-export function reduceSessionRows(messages: SessionMessageInfo[], inputs = new Set<string>()) {
+export function reduceSessionRows(
+  messages: SessionMessageInfo[],
+  inputs = new Set<string>(),
+  turnTokens = false,
+) {
   const isInput = (message: SessionMessageInfo) => inputs.has(message.id)
   const pendingCompactions = messages.filter((message) => message.type === "compaction" && message.status === "running")
   const pending = new Set([...pendingCompactions.map((message) => message.id), ...inputs])
-  const steps: string[] = []
-  let previousCacheRead: number | undefined
-  let turnPreviousCacheRead: number | undefined
-  let measured = false
+  const usage = turnTokens
+    ? { steps: [] as SessionMessageAssistant[], previousTurnCacheRead: undefined as number | undefined }
+    : undefined
   return [
     ...messages.filter((message) => !pending.has(message.id)),
     ...pendingCompactions,
@@ -296,12 +293,7 @@ export function reduceSessionRows(messages: SessionMessageInfo[], inputs = new S
       rows.push({ type: "message", messageID: message.id })
       return rows
     }
-    if (steps.length === 0) turnPreviousCacheRead = previousCacheRead
-    steps.push(message.id)
-    if (message.tokens && tokenTotal(message.tokens) > 0) {
-      previousCacheRead = message.tokens.cache.read
-      measured = true
-    }
+    usage?.steps.push(message)
     const ordinals = { text: 0, reasoning: 0 }
     message.content.forEach((part) => {
       const partID = part.type === "tool" ? part.id : `${part.type}:${ordinals[part.type]++}`
@@ -313,21 +305,31 @@ export function reduceSessionRows(messages: SessionMessageInfo[], inputs = new S
       completePrevious(rows)
       rows.push({ type: "assistant-footer", messageID: message.id })
     }
-    if (terminal) {
-      if (measured)
+    if (terminal && usage) {
+      const stepsWithUsage = usage.steps.filter(hasTokenUsage)
+      const last = stepsWithUsage.at(-1)
+      if (last) {
         rows.push({
           type: "turn-usage",
-          messageIDs: [...steps],
-          ...(turnPreviousCacheRead === undefined ? {} : { previousCacheRead: turnPreviousCacheRead }),
+          messageIDs: stepsWithUsage.map((step) => step.id),
+          ...(usage.previousTurnCacheRead === undefined
+            ? {}
+            : { previousCacheRead: usage.previousTurnCacheRead }),
         })
-      steps.length = 0
-      turnPreviousCacheRead = undefined
-      measured = false
+        usage.previousTurnCacheRead = last.tokens.cache.read
+      }
+      usage.steps.length = 0
     }
     return rows
   }, [])
 }
 
+function hasTokenUsage(
+  message: SessionMessageAssistant,
+): message is SessionMessageAssistant & { tokens: NonNullable<SessionMessageAssistant["tokens"]> } {
+  return message.tokens !== undefined && tokenTotal(message.tokens) > 0
+}
+
 function tokenTotal(tokens: NonNullable<SessionMessageAssistant["tokens"]>) {
   return tokens.input + tokens.output + tokens.reasoning + tokens.cache.read + tokens.cache.write
 }

packages/tui/test/cli/tui/data.test.tsx

@@ -5,12 +5,14 @@ import type { OpenCodeEvent } from "@opencode-ai/client"
 import { SessionMessage } from "@opencode-ai/core/session/message"
 import { EventV2 } from "@opencode-ai/core/event"
 import { createEffect, onMount, type ParentProps } from "solid-js"
+import { ConfigProvider } from "../../../src/config"
 import { ClientProvider, useClient } from "../../../src/context/client"
 import { DataProvider as DataProviderBase, useData } from "../../../src/context/data"
 import { LocationProvider, useLocation } from "../../../src/context/location"
 import { createSessionRows, type SessionRow } from "../../../src/routes/session/rows"
 import { createApi, createEventStream, createFetch, directory, json } from "../../fixture/tui-client"
 import { TestTuiContexts } from "../../fixture/tui-environment"
+import { createTuiResolvedConfig } from "../../fixture/tui-runtime"
 
 const formFields = [{ key: "authorization", type: "external", url: "https://example.com" }] satisfies [
   {
@@ -32,14 +34,18 @@ function emitEvent(events: ReturnType<typeof createEventStream>, event: OpenCode
   events.emit({ ...event, location: { directory } })
 }
 
+const config = createTuiResolvedConfig()
+
 function DataProvider(props: ParentProps) {
   return (
-    <DataProviderBase>
-      <LocationProvider>
-        <SyncLocation />
-        {props.children}
-      </LocationProvider>
-    </DataProviderBase>
+    <ConfigProvider config={config}>
+      <DataProviderBase>
+        <LocationProvider>
+          <SyncLocation />
+          {props.children}
+        </LocationProvider>
+      </DataProviderBase>
+    </ConfigProvider>
   )
 }