Refine session usage tracking

packages/ui/src/components/message-item.tsx

@@ -1,6 +1,7 @@
-import { For, Show } from "solid-js"
+import { For, Show, createMemo } from "solid-js"
 import type { Message, SDKPart, MessageInfo, ClientPart } from "../types/message"
 import { partHasRenderableText } from "../types/message"
+import { formatTokenTotal } from "../lib/formatters"
 import MessagePart from "./message-part"
 
 interface MessageItemProps {
@@ -138,6 +139,41 @@ export default function MessageItem(props: MessageItemProps) {
     isUser()
       ? "message-item-base bg-[var(--message-user-bg)] border-l-4 border-[var(--message-user-border)]"
       : "message-item-base assistant-message bg-[var(--message-assistant-bg)] border-l-4 border-[var(--message-assistant-border)]"
+
+  const statChipClass =
+    "inline-flex items-center gap-1 rounded-full border border-[var(--border-base)] px-2 py-0.5 text-[10px]"
+  const statLabelClass = "uppercase text-[9px] tracking-wide text-[var(--text-muted)]"
+  const statValueClass = "font-semibold text-[var(--text-primary)]"
+
+  const usageStats = createMemo(() => {
+    const info = props.messageInfo
+    if (!info || info.role !== "assistant" || !info.tokens) {
+      return null
+    }
+
+    const tokens = info.tokens
+    const input = tokens.input ?? 0
+    const output = tokens.output ?? 0
+    const reasoning = tokens.reasoning ?? 0
+    if (input === 0 && output === 0 && reasoning === 0) {
+      return null
+    }
+
+    return {
+      input,
+      output,
+      reasoning,
+      cacheRead: tokens.cache?.read ?? 0,
+      cacheWrite: tokens.cache?.write ?? 0,
+      cost: info.cost ?? 0,
+    }
+  })
+
+  const formatCostValue = (value: number) => {
+    if (!value) return "$0.00"
+    if (value < 0.01) return `$${value.toPrecision(2)}`
+    return `$${value.toFixed(2)}`
+  }
  
   const agentIdentifier = () => {
     if (isUser()) return ""
@@ -193,6 +229,36 @@ export default function MessageItem(props: MessageItemProps) {
               Fork
             </button>
           </Show>
+          <Show when={usageStats()}>
+            {(usage) => (
+              <div class="flex flex-wrap items-center gap-1 text-[10px] text-[var(--text-muted)]">
+                <div class={statChipClass}>
+                  <span class={statLabelClass}>Input</span>
+                  <span class={statValueClass}>{formatTokenTotal(usage().input)}</span>
+                </div>
+                <div class={statChipClass}>
+                  <span class={statLabelClass}>Output</span>
+                  <span class={statValueClass}>{formatTokenTotal(usage().output)}</span>
+                </div>
+                <div class={statChipClass}>
+                  <span class={statLabelClass}>Reasoning</span>
+                  <span class={statValueClass}>{formatTokenTotal(usage().reasoning)}</span>
+                </div>
+                <div class={statChipClass}>
+                  <span class={statLabelClass}>Cache Read</span>
+                  <span class={statValueClass}>{formatTokenTotal(usage().cacheRead)}</span>
+                </div>
+                <div class={statChipClass}>
+                  <span class={statLabelClass}>Cache Write</span>
+                  <span class={statValueClass}>{formatTokenTotal(usage().cacheWrite)}</span>
+                </div>
+                <div class={statChipClass}>
+                  <span class={statLabelClass}>Cost</span>
+                  <span class={statValueClass}>{formatCostValue(usage().cost)}</span>
+                </div>
+              </div>
+            )}
+          </Show>
           <span class="text-[11px] text-[var(--text-muted)]">{timestamp()}</span>
         </div>
       </div>

packages/ui/src/components/message-stream.tsx

@@ -32,6 +32,7 @@ import { sseManager } from "../lib/sse-manager"
 import Kbd from "./kbd"
 import { useConfig } from "../stores/preferences"
 import { getSessionInfo, computeDisplayParts, sessions, setActiveSession, setActiveParentSession } from "../stores/sessions"
+import { formatTokenTotal } from "../lib/formatters"
 import { setActiveInstanceId } from "../stores/instances"
 import { showCommandPalette } from "../stores/command-palette"
 
@@ -72,26 +73,9 @@ function navigateToTaskSession(location: TaskSessionLocation) {
   }
 }
 
-// Format tokens like TUI (e.g., "110K", "1.2M")
+// Format tokens like session sidebar (comma-separated totals)
 function formatTokens(tokens: number): string {
-  if (tokens >= 1000000) {
-    return `${(tokens / 1000000).toFixed(1)}M`
-  } else if (tokens >= 1000) {
-    return `${(tokens / 1000).toFixed(0)}K`
-  }
-  return tokens.toString()
-}
-
-// Format session info for the session view header
-function formatSessionInfo(usageTokens: number, contextWindow: number, usagePercent: number | null): string {
-  if (contextWindow > 0) {
-    const windowStr = formatTokens(contextWindow)
-    const usageStr = formatTokens(usageTokens)
-    const percent = usagePercent ?? Math.min(100, Math.max(0, Math.round((usageTokens / contextWindow) * 100)))
-    return `${usageStr} of ${windowStr} (${percent}%)`
-  }
-
-  return formatTokens(usageTokens)
+  return formatTokenTotal(tokens)
 }
 
 interface MessageStreamProps {
@@ -206,18 +190,27 @@ export default function MessageStream(props: MessageStreamProps) {
 
   const sessionInfo = createMemo(() =>
     getSessionInfo(props.instanceId, props.sessionId) ?? {
-      tokens: 0,
       cost: 0,
       contextWindow: 0,
       isSubscriptionModel: false,
-      contextUsageTokens: 0,
-      contextUsagePercent: null,
+      inputTokens: 0,
+      outputTokens: 0,
+      reasoningTokens: 0,
+      actualUsageTokens: 0,
+      modelOutputLimit: 0,
+      contextAvailableTokens: null,
     },
   )
 
-  const formattedSessionInfo = createMemo(() => {
+  const tokenStats = createMemo(() => {
     const info = sessionInfo()
-    return formatSessionInfo(info.contextUsageTokens, info.contextWindow, info.contextUsagePercent)
+    return {
+      input: info.inputTokens ?? 0,
+      output: info.outputTokens ?? 0,
+      cost: info.cost ?? 0,
+      used: info.actualUsageTokens ?? 0,
+      avail: info.contextAvailableTokens,
+    }
   })
 
   function isNearBottom(element: HTMLDivElement, offset = SCROLL_OFFSET) {
@@ -552,10 +545,20 @@ export default function MessageStream(props: MessageStreamProps) {
 
   return (
     <div class="message-stream-container">
-      <div class="connection-status">
-        <div class="connection-status-text connection-status-info flex items-center gap-2 text-sm font-medium">
-          <span>{formattedSessionInfo()}</span>
-        </div>
+        <div class="connection-status">
+          <div class="connection-status-text connection-status-info flex flex-wrap items-center gap-2 text-sm font-medium">
+            <div class="inline-flex items-center gap-1 rounded-full border border-base px-2 py-0.5 text-xs text-primary">
+              <span class="uppercase text-[10px] tracking-wide text-primary/70">Used</span>
+              <span class="font-semibold text-primary">{formatTokens(sessionInfo().actualUsageTokens ?? 0)}</span>
+            </div>
+            <div class="inline-flex items-center gap-1 rounded-full border border-base px-2 py-0.5 text-xs text-primary">
+              <span class="uppercase text-[10px] tracking-wide text-primary/70">Avail</span>
+              <span class="font-semibold text-primary">
+                {sessionInfo().contextAvailableTokens !== null ? formatTokens(sessionInfo().contextAvailableTokens ?? 0) : "--"}
+              </span>
+            </div>
+          </div>
+
         <div class="connection-status-text connection-status-shortcut">
           <div class="connection-status-shortcut-action">
             <button

packages/ui/src/components/session/context-usage-panel.tsx

@@ -7,54 +7,71 @@ interface ContextUsagePanelProps {
   sessionId: string
 }
 
+const chipClass = "inline-flex items-center gap-1 rounded-full border border-base px-2 py-0.5 text-xs text-primary"
+const chipLabelClass = "uppercase text-[10px] tracking-wide text-primary/70"
+const headingClass = "text-xs font-semibold text-primary/70 uppercase tracking-wide"
+
 const ContextUsagePanel: Component<ContextUsagePanelProps> = (props) => {
   const info = createMemo(
     () =>
       getSessionInfo(props.instanceId, props.sessionId) ?? {
-        tokens: 0,
         cost: 0,
         contextWindow: 0,
         isSubscriptionModel: false,
-        contextUsageTokens: 0,
-        contextUsagePercent: null,
+        inputTokens: 0,
+        outputTokens: 0,
+        reasoningTokens: 0,
+        actualUsageTokens: 0,
+        modelOutputLimit: 0,
+        contextAvailableTokens: null,
       },
   )
 
-  const tokens = createMemo(() => info().tokens)
-  const contextUsageTokens = createMemo(() => info().contextUsageTokens ?? 0)
-  const contextWindow = createMemo(() => info().contextWindow)
-  const contextUsagePercent = createMemo(() => info().contextUsagePercent)
-
-  const costLabel = createMemo(() => {
-    if (info().isSubscriptionModel || info().cost <= 0) return "Included in plan"
-    return `$${info().cost.toFixed(2)} spent`
+  const inputTokens = createMemo(() => info().inputTokens ?? 0)
+  const outputTokens = createMemo(() => info().outputTokens ?? 0)
+  const actualUsageTokens = createMemo(() => info().actualUsageTokens ?? 0)
+  const availableTokens = createMemo(() => info().contextAvailableTokens)
+  const outputLimit = createMemo(() => info().modelOutputLimit ?? 0)
+  const costValue = createMemo(() => {
+    const value = info().isSubscriptionModel ? 0 : info().cost
+    return value > 0 ? value : 0
   })
 
+  const formatTokenValue = (value: number | null | undefined) => {
+    if (value === null || value === undefined) return "--"
+    return formatTokenTotal(value)
+  }
+
+  const costDisplay = createMemo(() => `$${costValue().toFixed(2)}`)
+
   return (
-    <div class="session-context-panel border-r border-base border-b px-3 py-3">
-      <div class="flex items-center justify-between gap-4">
-        <div>
-          <div class="text-xs font-semibold text-primary/70 uppercase tracking-wide">Tokens (last call)</div>
-          <div class="text-lg font-semibold text-primary">{formatTokenTotal(tokens())}</div>
+    <div class="session-context-panel border-r border-base border-b px-3 py-3 space-y-3">
+      <div class="flex flex-wrap items-center gap-2 text-xs text-primary/90">
+        <div class={headingClass}>Tokens</div>
+        <div class={chipClass}>
+          <span class={chipLabelClass}>Input</span>
+          <span class="font-semibold text-primary">{formatTokenTotal(inputTokens())}</span>
         </div>
-        <div class="text-xs text-primary/70 text-right leading-tight">{costLabel()}</div>
-      </div>
-      <div class="mt-4">
-        <div class="flex items-center justify-between mb-1">
-          <div class="text-xs font-semibold text-primary/70 uppercase tracking-wide">Context window usage</div>
-          <div class="text-sm font-medium text-primary">{contextUsagePercent() !== null ? `${contextUsagePercent()}%` : "--"}</div>
+        <div class={chipClass}>
+          <span class={chipLabelClass}>Output</span>
+          <span class="font-semibold text-primary">{formatTokenTotal(outputTokens())}</span>
         </div>
-        <div class="text-sm text-primary/90">
-          {contextWindow()
-            ? `${formatTokenTotal(contextUsageTokens())} of ${formatTokenTotal(contextWindow())}`
-            : "Window size unavailable"}
+        <div class={chipClass}>
+          <span class={chipLabelClass}>Cost</span>
+          <span class="font-semibold text-primary">{costDisplay()}</span>
         </div>
       </div>
-      <div class="mt-3 h-1.5 rounded-full bg-base relative overflow-hidden">
-        <div
-          class="absolute inset-y-0 left-0 rounded-full bg-accent-primary transition-[width]"
-          style={{ width: contextUsagePercent() === null ? "0%" : `${contextUsagePercent()}%` }}
-        />
+
+      <div class="flex flex-wrap items-center gap-2 text-xs text-primary/90">
+        <div class={headingClass}>Context</div>
+        <div class={chipClass}>
+          <span class={chipLabelClass}>Used</span>
+          <span class="font-semibold text-primary">{formatTokenTotal(actualUsageTokens())}</span>
+        </div>
+        <div class={chipClass}>
+          <span class={chipLabelClass}>Avail</span>
+          <span class="font-semibold text-primary">{formatTokenValue(availableTokens())}</span>
+        </div>
       </div>
     </div>
   )