feat(tui): add token and cost info to session header

1 files changed, 101 insertions, 2 deletions

diff --git a/packages/tui/internal/components/chat/messages.go b/packages/tui/internal/components/chat/messages.go
index 573acdfcc..7ecd9b21f 100644
--- a/packages/tui/internal/components/chat/messages.go
+++ b/packages/tui/internal/components/chat/messages.go
@@ -1,6 +1,7 @@
 package chat
 
 import (
+	"fmt"
 	"strings"
 
 	"github.com/charmbracelet/bubbles/v2/viewport"
@@ -371,11 +372,65 @@ func (m *messagesComponent) header(width int) string {
 		headerLines,
 		util.ToMarkdown("# "+m.app.Session.Title, width-6, t.Background()),
 	)
+
+	share := ""
 	if m.app.Session.Share.URL != "" {
-		headerLines = append(headerLines, muted(m.app.Session.Share.URL+"  /unshare"))
+		share = muted(m.app.Session.Share.URL + "  /unshare")
 	} else {
-		headerLines = append(headerLines, base("/share")+muted(" to create a shareable link"))
+		share = base("/share") + muted(" to create a shareable link")
+	}
+
+	sessionInfo := ""
+	tokens := float64(0)
+	cost := float64(0)
+	contextWindow := m.app.Model.Limit.Context
+
+	for _, message := range m.app.Messages {
+		if assistant, ok := message.(opencode.AssistantMessage); ok {
+			cost += assistant.Cost
+			usage := assistant.Tokens
+			if usage.Output > 0 {
+				if assistant.Summary {
+					tokens = usage.Output
+					continue
+				}
+				tokens = (usage.Input +
+					usage.Cache.Write +
+					usage.Cache.Read +
+					usage.Output +
+					usage.Reasoning)
+			}
+		}
 	}
+
+	// Check if current model is a subscription model (cost is 0 for both input and output)
+	isSubscriptionModel := m.app.Model != nil &&
+		m.app.Model.Cost.Input == 0 && m.app.Model.Cost.Output == 0
+
+	sessionInfo = styles.NewStyle().
+		Foreground(t.TextMuted()).
+		Background(t.Background()).
+		Render(formatTokensAndCost(tokens, contextWindow, cost, isSubscriptionModel))
+
+	background := t.Background()
+	share = layout.Render(
+		layout.FlexOptions{
+			Background: &background,
+			Direction:  layout.Row,
+			Justify:    layout.JustifySpaceBetween,
+			Align:      layout.AlignStretch,
+			Width:      width - 6,
+		},
+		layout.FlexItem{
+			View: share,
+		},
+		layout.FlexItem{
+			View: sessionInfo,
+		},
+	)
+
+	headerLines = append(headerLines, share)
+
 	header := strings.Join(headerLines, "\n")
 
 	header = styles.NewStyle().
@@ -393,6 +448,50 @@ func (m *messagesComponent) header(width int) string {
 	return "\n" + header + "\n"
 }
 
+func formatTokensAndCost(
+	tokens float64,
+	contextWindow float64,
+	cost float64,
+	isSubscriptionModel bool,
+) string {
+	// Format tokens in human-readable format (e.g., 110K, 1.2M)
+	var formattedTokens string
+	switch {
+	case tokens >= 1_000_000:
+		formattedTokens = fmt.Sprintf("%.1fM", float64(tokens)/1_000_000)
+	case tokens >= 1_000:
+		formattedTokens = fmt.Sprintf("%.1fK", float64(tokens)/1_000)
+	default:
+		formattedTokens = fmt.Sprintf("%d", int(tokens))
+	}
+
+	// Remove .0 suffix if present
+	if strings.HasSuffix(formattedTokens, ".0K") {
+		formattedTokens = strings.Replace(formattedTokens, ".0K", "K", 1)
+	}
+	if strings.HasSuffix(formattedTokens, ".0M") {
+		formattedTokens = strings.Replace(formattedTokens, ".0M", "M", 1)
+	}
+
+	percentage := (float64(tokens) / float64(contextWindow)) * 100
+
+	if isSubscriptionModel {
+		return fmt.Sprintf(
+			"%s/%d%%",
+			formattedTokens,
+			int(percentage),
+		)
+	}
+
+	formattedCost := fmt.Sprintf("$%.2f", cost)
+	return fmt.Sprintf(
+		"%s/%d%% (%s)",
+		formattedTokens,
+		int(percentage),
+		formattedCost,
+	)
+}
+
 func (m *messagesComponent) View(width, height int) string {
 	t := theme.CurrentTheme()
 	if m.rendering {