Procházet zdrojové kódy

feat(tui): add turn token usage diagnostics (#38398)

James Long před 3 týdny
rodič
revize
5b1321a8ca

+ 12 - 1
packages/tui/src/component/devtools-bar.tsx

@@ -59,6 +59,7 @@ export function DevToolsBar() {
   const canSwitchMode = () => supports(nextMode())
   const runtime = createMemo(() => runtimeStatus(frontendSamples()))
   const timing = () => config.data.debug?.timing ?? false
+  const turnTokens = () => config.data.debug?.turn_tokens ?? false
 
   const offEscape = keymap.intercept(
     "key",
@@ -352,6 +353,16 @@ export function DevToolsBar() {
               >
                 {timing() ? "[x]" : "[ ]"} Time to first draw
               </Action>
+              <Action
+                onClick={() =>
+                  void config.update((draft) => {
+                    draft.debug = { ...draft.debug, turn_tokens: !turnTokens() }
+                  })
+                }
+                hoverBackground
+              >
+                {turnTokens() ? "[x]" : "[ ]"} Turn token usage
+              </Action>
             </box>
             <For each={groups()}>
               {(group) => (
@@ -403,7 +414,7 @@ function PanelBox(props: ParentProps) {
       position="absolute"
       zIndex={2600}
       bottom={1}
-      left={0}
+      left={-1}
       width={42}
       paddingLeft={2}
       paddingRight={2}

+ 0 - 8
packages/tui/src/component/dialog-config.tsx

@@ -222,14 +222,6 @@ const settings: Setting[] = [
     values: [false, true],
     labels: ["off", "on"],
   },
-  {
-    title: "DevTools: Timing",
-    category: "Debug",
-    path: ["debug", "timing"],
-    default: true,
-    values: [false, true],
-    labels: ["off", "on"],
-  },
 ]
 
 export function DialogConfig() {

+ 1 - 0
packages/tui/src/config/index.tsx

@@ -159,6 +159,7 @@ export const Info = Schema.Struct({
     Schema.Struct({
       devtools: Schema.optional(Schema.Boolean).annotate({ description: "Show the DevTools debug bar" }),
       timing: Schema.optional(Schema.Boolean).annotate({ description: "Show time-to-first-draw diagnostics" }),
+      turn_tokens: Schema.optional(Schema.Boolean).annotate({ description: "Show per-turn token usage diagnostics" }),
     }),
   ).annotate({ description: "Debugging settings" }),
   animations: Schema.optional(Schema.Boolean).annotate({ description: "Enable interface animations" }),

+ 112 - 2
packages/tui/src/routes/session/index.tsx

@@ -1029,11 +1029,23 @@ export function Session() {
   )
 }
 
-function SessionRowView(props: {
+type SessionRowViewProps = {
   row: SessionRow
   message: (messageID: string) => SessionMessageInfo | undefined
   boundaryID?: string
-}) {
+}
+
+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>
@@ -1072,11 +1084,109 @@ function SessionRowView(props: {
             </Show>
           )}
         </Match>
+        <Match when={props.row.type === "turn-usage" ? props.row : undefined}>
+          {(row) => (
+            <TurnTokenUsage
+              messageIDs={row().messageIDs}
+              previousCacheRead={row().previousCacheRead}
+              message={props.message}
+            />
+          )}
+        </Match>
       </Switch>
     </box>
   )
 }
 
+function TurnTokenUsage(props: {
+  messageIDs: string[]
+  previousCacheRead?: number
+  message: (messageID: string) => SessionMessageInfo | undefined
+}) {
+  const config = useConfig()
+  const { themeV2 } = useTheme()
+  const steps = createMemo(() => {
+    let previousCacheRead = props.previousCacheRead
+    return props.messageIDs.flatMap((messageID) => {
+      const message = props.message(messageID)
+      if (message?.type !== "assistant" || !message.tokens) return []
+      const total =
+        message.tokens.input +
+        message.tokens.output +
+        message.tokens.reasoning +
+        message.tokens.cache.read +
+        message.tokens.cache.write
+      if (total === 0) return []
+      const newTokens = total - message.tokens.cache.read
+      const cacheBust =
+        previousCacheRead !== undefined && message.tokens.cache.read < previousCacheRead
+          ? previousCacheRead - message.tokens.cache.read
+          : undefined
+      previousCacheRead = message.tokens.cache.read
+      return [
+        {
+          finish: message.finish === "tool-calls" ? "tool-call" : (message.finish ?? "unknown"),
+          newTokens,
+          cached: message.tokens.cache.read,
+          total,
+          cacheBust,
+        },
+      ]
+    })
+  })
+  const columns = createMemo(() => ({
+    step: Math.max("Step".length, ...steps().map((item) => item.finish.length)),
+    newTokens: Math.max("New".length, ...steps().map((item) => item.newTokens.toLocaleString().length)),
+    cached: Math.max("Cached".length, ...steps().map((item) => item.cached.toLocaleString().length)),
+    total: Math.max("Total".length, ...steps().map((item) => item.total.toLocaleString().length)),
+  }))
+  return (
+    <Show when={config.data.debug?.turn_tokens === true && steps().length > 0}>
+      <box paddingLeft={3} flexDirection="column">
+        <box flexDirection="row">
+          <text width={INLINE_TOOL_ICON_WIDTH} fg={themeV2.text.subdued}>
+            ◈
+          </text>
+          <text fg={themeV2.text.subdued} attributes={TextAttributes.BOLD}>
+            Tokens
+          </text>
+        </box>
+        <box paddingLeft={INLINE_TOOL_ICON_WIDTH}>
+          <text fg={themeV2.text.subdued} attributes={TextAttributes.ITALIC}>
+            {"Step".padEnd(columns().step + 2)}
+            {"New".padStart(columns().newTokens)}
+            {"  "}
+            {"Cached".padStart(columns().cached)}
+            {"  "}
+            {"Total".padStart(columns().total)}
+          </text>
+        </box>
+        <For each={steps()}>
+          {(item) => (
+            <box paddingLeft={INLINE_TOOL_ICON_WIDTH} flexDirection="column">
+              <text fg={themeV2.text.subdued}>
+                {item.finish.padEnd(columns().step + 2)}
+                <span style={{ attributes: TextAttributes.BOLD }}>
+                  {item.newTokens.toLocaleString().padStart(columns().newTokens)}
+                </span>
+                {"  "}
+                {item.cached.toLocaleString().padStart(columns().cached)}
+                {"  "}
+                {item.total.toLocaleString().padStart(columns().total)}
+              </text>
+              <Show when={item.cacheBust !== undefined}>
+                <text fg={themeV2.text.feedback.error.default}>
+                  ! Cache bust: {item.cacheBust?.toLocaleString()} fewer cached tokens than the previous step
+                </text>
+              </Show>
+            </box>
+          )}
+        </For>
+      </box>
+    </Show>
+  )
+}
+
 function BackgroundToolHint(props: { messages: SessionMessageInfo[] }) {
   const { themeV2 } = useTheme()
   const shortcut = Keymap.useShortcut("session.background")

+ 50 - 1
packages/tui/src/routes/session/rows.ts

@@ -27,6 +27,7 @@ export type SessionRow =
       completed: boolean
     }
   | { type: "assistant-footer"; messageID: string }
+  | { type: "turn-usage"; messageIDs: string[]; previousCacheRead?: number }
 
 export function createSessionRows(sessionID: Accessor<string>) {
   const data = useData()
@@ -127,6 +128,26 @@ 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())),
+    ),
+  )
+
   const appendMessage = (messageID: string) =>
     setRows(
       produce((draft) => {
@@ -260,6 +281,10 @@ export function reduceSessionRows(messages: SessionMessageInfo[], inputs = new S
   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
   return [
     ...messages.filter((message) => !pending.has(message.id)),
     ...pendingCompactions,
@@ -271,20 +296,42 @@ 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
+    }
     const ordinals = { text: 0, reasoning: 0 }
     message.content.forEach((part) => {
       const partID = part.type === "tool" ? part.id : `${part.type}:${ordinals[part.type]++}`
       if ((part.type === "text" || part.type === "reasoning") && !part.text.trim()) return
       append(rows, { messageID: message.id, partID }, part)
     })
-    if ((message.finish && !["tool-calls", "unknown"].includes(message.finish)) || message.error || message.retry) {
+    const terminal = (message.finish && !["tool-calls", "unknown"].includes(message.finish)) || message.error
+    if (terminal || message.retry) {
       completePrevious(rows)
       rows.push({ type: "assistant-footer", messageID: message.id })
     }
+    if (terminal) {
+      if (measured)
+        rows.push({
+          type: "turn-usage",
+          messageIDs: [...steps],
+          ...(turnPreviousCacheRead === undefined ? {} : { previousCacheRead: turnPreviousCacheRead }),
+        })
+      steps.length = 0
+      turnPreviousCacheRead = undefined
+      measured = false
+    }
     return rows
   }, [])
 }
 
+function tokenTotal(tokens: NonNullable<SessionMessageAssistant["tokens"]>) {
+  return tokens.input + tokens.output + tokens.reasoning + tokens.cache.read + tokens.cache.write
+}
+
 export function messageBoundaryIDs(rows: SessionRow[], messages: SessionMessageInfo[]) {
   const byID = new Map(messages.map((message) => [message.id, message]))
   const seen = new Set<string>()
@@ -309,6 +356,8 @@ function rowBoundaryMessageID(row: SessionRow, messages: Map<string, SessionMess
         ? row.refs[0]?.messageID
         : row.type === "assistant-footer"
           ? row.messageID
+          : row.type === "turn-usage"
+            ? row.messageIDs[0]
           : undefined
   if (!messageID) return undefined
   const message = messages.get(messageID)