From 0ef57df3bc681e9e362b26d57cb00256132efd5b Mon Sep 17 00:00:00 2001
From: Shantur Rathore <i@shantur.com>
Date: Wed, 8 Apr 2026 22:19:10 +0100
Subject: [PATCH] fix(ui): show token stats and simplify context window
 calculation

- Track messageInfoVersion in cache signature to rebuild when tokens arrive via SSE
- Read tokens from step-finish part directly (embedded in SSE events)
- Simplify available tokens to show full context window when no explicit input limit
---
 packages/ui/src/components/message-block.tsx  | 20 +++++++++++++------
 .../ui/src/stores/message-v2/session-info.ts  | 17 +++++-----------
 2 files changed, 19 insertions(+), 18 deletions(-)

diff --git a/packages/ui/src/components/message-block.tsx b/packages/ui/src/components/message-block.tsx
index 211127a9..4e856411 100644
--- a/packages/ui/src/components/message-block.tsx
+++ b/packages/ui/src/components/message-block.tsx
@@ -629,13 +629,12 @@ export default function MessageBlock(props: MessageBlockProps) {
     const lastAssistantIdx = props.lastAssistantIndex()
     const isQueued = current.role === "user" && (lastAssistantIdx === -1 || index > lastAssistantIdx)
 
-    // Intentionally untracked: messageInfoVersion updates should not trigger
-    // a full message block rebuild; record revision is the invalidation key.
-    const info = untrack(messageInfo)
+    const messageInfoVersion = props.store().state.messageInfoVersion[current.id] ?? 0
 
     const cacheSignature = [
       current.id,
       current.revision,
+      messageInfoVersion,
       isQueued ? 1 : 0,
       props.showThinking() ? 1 : 0,
       props.thinkingDefaultExpanded() ? 1 : 0,
@@ -647,6 +646,9 @@ export default function MessageBlock(props: MessageBlockProps) {
       return cachedBlock.block
     }
 
+    // Only capture info after cache check fails - ensures fresh data on version bump
+    const info = untrack(messageInfo)
+
     const { orderedParts } = buildRecordDisplayData(props.instanceId, current)
     const items: MessageBlockItem[] = []
     const blockContentKeys: string[] = []
@@ -1108,17 +1110,23 @@ function StepCard(props: StepCardProps) {
       return null
     }
     const info = props.messageInfo
-    if (!info || info.role !== "assistant" || !info.tokens) {
+    const part = props.part as any
+    
+    // step-finish parts have tokens embedded; also check messageInfo
+    const partTokens = part?.tokens
+    const infoTokens = info && info.role === "assistant" ? info.tokens : undefined
+    const tokens = partTokens ?? infoTokens
+    if (!tokens) {
       return null
     }
-    const tokens = info.tokens
+    
     return {
       input: tokens.input ?? 0,
       output: tokens.output ?? 0,
       reasoning: tokens.reasoning ?? 0,
       cacheRead: tokens.cache?.read ?? 0,
       cacheWrite: tokens.cache?.write ?? 0,
-      cost: info.cost ?? 0,
+      cost: (part?.cost ?? (info && info.role === "assistant" ? info.cost : 0)) ?? 0,
     }
   }
 
diff --git a/packages/ui/src/stores/message-v2/session-info.ts b/packages/ui/src/stores/message-v2/session-info.ts
index a0970ebf..50bbf697 100644
--- a/packages/ui/src/stores/message-v2/session-info.ts
+++ b/packages/ui/src/stores/message-v2/session-info.ts
@@ -116,18 +116,11 @@ export function updateSessionInfo(instanceId: string, sessionId: string): void {
     // Prefer explicit input limits when provided by the API.
     // This is used by the UI "Avail" chip.
     contextAvailableTokens = modelInputLimit
-  }
-
-  if (!contextAvailableFromPrevious && contextAvailableTokens === null) {
-    if (contextWindow > 0) {
-      if (latestHasContextUsage && actualUsageTokens > 0) {
-        contextAvailableTokens = Math.max(contextWindow - (actualUsageTokens + outputBudget), 0)
-      } else {
-        contextAvailableTokens = contextWindow
-      }
-    } else {
-      contextAvailableTokens = null
-    }
+  } else if (contextWindow > 0) {
+    // When no explicit input limit, show full context window capacity.
+    contextAvailableTokens = contextWindow
+  } else {
+    contextAvailableTokens = null
   }
 
   setSessionInfoByInstance((prev) => {