From 2736e487cf511a686daf691d9135f9491d4ac865 Mon Sep 17 00:00:00 2001 From: Greg Slepak Date: Tue, 12 May 2026 17:34:08 -0700 Subject: [PATCH] fix(ui): mark estimated context usage --- internal/agent/agent.go | 5 +++++ internal/agent/usage_fallback_test.go | 2 ++ internal/session/session.go | 3 +++ internal/ui/common/elements.go | 17 +++++++++----- internal/ui/common/elements_test.go | 32 +++++++++++++++++++++++++++ internal/ui/model/header.go | 6 ++++- internal/ui/model/sidebar.go | 7 +++--- internal/ui/styles/quickstyle.go | 1 + internal/ui/styles/styles.go | 21 +++++++++--------- 9 files changed, 74 insertions(+), 20 deletions(-) create mode 100644 internal/ui/common/elements_test.go diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 15b0c1af161a7e5631ba00e55e988a19aa28fe0e..407604978bbd54d6e38485a6ac0af6f6315edaf4 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -762,6 +762,7 @@ func (a *sessionAgent) Summarize(ctx context.Context, sessionID string, opts fan currentSession.SummaryMessageID = summaryMessage.ID currentSession.CompletionTokens = summaryCompletionTokens(usage, summaryMessage) currentSession.PromptTokens = 0 + currentSession.EstimatedUsage = usageIsZero(usage) _, err = a.sessions.Save(genCtx, currentSession) if err != nil { return err @@ -1139,6 +1140,10 @@ func (a *sessionAgent) openrouterCost(metadata fantasy.ProviderMetadata) *float6 } func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session, usage fantasy.Usage, overrideCost *float64, estimated bool) { + if !usageIsZero(usage) { + session.EstimatedUsage = estimated + } + modelConfig := model.CatwalkCfg cost := modelConfig.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) + modelConfig.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) + diff --git a/internal/agent/usage_fallback_test.go b/internal/agent/usage_fallback_test.go index 1f2442aef7d4a1623419f70c81cd494ccc88ef69..3b6fc19b943fbb1cf636782c051e91a6ff64a5df 100644 --- a/internal/agent/usage_fallback_test.go +++ b/internal/agent/usage_fallback_test.go @@ -212,6 +212,7 @@ func TestUpdateSessionUsageSkipsEstimatedCost(t *testing.T) { require.Equal(t, 1.25, currentSession.Cost) require.Equal(t, int64(1000), currentSession.PromptTokens) require.Equal(t, int64(2000), currentSession.CompletionTokens) + require.True(t, currentSession.EstimatedUsage) } func TestUpdateSessionUsageKeepsCountersForZeroUsage(t *testing.T) { @@ -334,4 +335,5 @@ func TestUpdateSessionUsageAddsProviderCost(t *testing.T) { require.Equal(t, 1.3, currentSession.Cost) require.Equal(t, int64(1000), currentSession.PromptTokens) require.Equal(t, int64(2000), currentSession.CompletionTokens) + require.False(t, currentSession.EstimatedUsage) } diff --git a/internal/session/session.go b/internal/session/session.go index 66bd9f4c9a12916d02c6d22ed7d51f81d74efdfd..6de6b9111d2f81fa49ae15e9ffaa9390f842d114 100644 --- a/internal/session/session.go +++ b/internal/session/session.go @@ -53,6 +53,7 @@ type Session struct { MessageCount int64 PromptTokens int64 CompletionTokens int64 + EstimatedUsage bool SummaryMessageID string Cost float64 Todos []Todo @@ -199,7 +200,9 @@ func (s *service) Save(ctx context.Context, session Session) (Session, error) { if err != nil { return Session{}, err } + estimatedUsage := session.EstimatedUsage session = s.fromDBItem(dbSession) + session.EstimatedUsage = estimatedUsage s.Publish(pubsub.UpdatedEvent, session) return session, nil } diff --git a/internal/ui/common/elements.go b/internal/ui/common/elements.go index 902c8a816247d97f0d9dd0ce7538b1bf130bc07d..de32645f2806b6c1790f10d4d0be2532cfdbea67 100644 --- a/internal/ui/common/elements.go +++ b/internal/ui/common/elements.go @@ -33,9 +33,10 @@ func FormatReasoningEffort(effort string) string { // ModelContextInfo contains token usage and cost information for a model. type ModelContextInfo struct { - ContextUsed int64 - ModelContext int64 - Cost float64 + ContextUsed int64 + ModelContext int64 + Cost float64 + EstimatedUsage bool } // ModelInfo renders model information including name, provider, reasoning @@ -74,7 +75,7 @@ func ModelInfo(t *styles.Styles, modelName, providerName, reasoningInfo string, } if context != nil { - formattedInfo := formatTokensAndCost(t, context.ContextUsed, context.ModelContext, context.Cost) + formattedInfo := formatTokensAndCost(t, context.ContextUsed, context.ModelContext, context.Cost, context.EstimatedUsage) parts = append(parts, lipgloss.NewStyle().PaddingLeft(2).Render(formattedInfo)) } @@ -92,7 +93,7 @@ func ModelInfo(t *styles.Styles, modelName, providerName, reasoningInfo string, // formatTokensAndCost formats token usage and cost with appropriate units // (K/M) and percentage of context window. -func formatTokensAndCost(t *styles.Styles, tokens, contextWindow int64, cost float64) string { +func formatTokensAndCost(t *styles.Styles, tokens, contextWindow int64, cost float64, estimated bool) string { var formattedTokens string switch { case tokens >= 1_000_000: @@ -115,7 +116,11 @@ func formatTokensAndCost(t *styles.Styles, tokens, contextWindow int64, cost flo formattedCost := t.ModelInfo.Cost.Render(fmt.Sprintf("$%.2f", cost)) formattedTokens = t.ModelInfo.TokenCount.Render(fmt.Sprintf("(%s)", formattedTokens)) - formattedPercentage := t.ModelInfo.TokenPercentage.Render(fmt.Sprintf("%d%%", int(percentage))) + percentageText := fmt.Sprintf("%d%%", int(percentage)) + if estimated { + percentageText = t.ModelInfo.EstimatedUsagePrefix.Render("~") + percentageText + } + formattedPercentage := t.ModelInfo.TokenPercentage.Render(percentageText) formattedTokens = fmt.Sprintf("%s %s", formattedPercentage, formattedTokens) if percentage > 80 { formattedTokens = fmt.Sprintf("%s %s", styles.LSPWarningIcon, formattedTokens) diff --git a/internal/ui/common/elements_test.go b/internal/ui/common/elements_test.go new file mode 100644 index 0000000000000000000000000000000000000000..801e4e52100fb6ed2d6960039da42b66fd609da6 --- /dev/null +++ b/internal/ui/common/elements_test.go @@ -0,0 +1,32 @@ +package common + +import ( + "testing" + + "github.com/charmbracelet/crush/internal/ui/styles" + "github.com/charmbracelet/x/ansi" + "github.com/stretchr/testify/require" +) + +func TestFormatTokensAndCostPrefixesEstimatedUsage(t *testing.T) { + t.Parallel() + + sty := styles.CharmtonePantera() + + actual := ansi.Strip(formatTokensAndCost(&sty, 120, 1000, 0, true)) + + require.Contains(t, actual, "~12%") + require.Contains(t, actual, "(120)") + require.Contains(t, actual, "$0.00") +} + +func TestFormatTokensAndCostOmitsEstimatedPrefix(t *testing.T) { + t.Parallel() + + sty := styles.CharmtonePantera() + + actual := ansi.Strip(formatTokensAndCost(&sty, 120, 1000, 0, false)) + + require.Contains(t, actual, "12%") + require.NotContains(t, actual, "~12%") +} diff --git a/internal/ui/model/header.go b/internal/ui/model/header.go index c7091a6835c48f06a2eb0d4a504d1cc11b79f09c..0d012f67ae9a6b8b1ef7d26d85d26b90a9e1c624 100644 --- a/internal/ui/model/header.go +++ b/internal/ui/model/header.go @@ -148,7 +148,11 @@ func renderHeaderDetails( model := com.Config().GetModelByType(agentCfg.Model) if model != nil && model.ContextWindow > 0 { percentage := (float64(session.CompletionTokens+session.PromptTokens) / float64(model.ContextWindow)) * 100 - formattedPercentage := t.Header.Percentage.Render(fmt.Sprintf("%d%%", int(percentage))) + percentageText := fmt.Sprintf("%d%%", int(percentage)) + if session.EstimatedUsage { + percentageText = t.ModelInfo.EstimatedUsagePrefix.Render("~") + percentageText + } + formattedPercentage := t.Header.Percentage.Render(percentageText) parts = append(parts, formattedPercentage) } diff --git a/internal/ui/model/sidebar.go b/internal/ui/model/sidebar.go index 2f548628f24168bbe5f7d4d8390fc24e472f9559..e98ef33423cff6147afe24ee97d6a18fe0c57ced 100644 --- a/internal/ui/model/sidebar.go +++ b/internal/ui/model/sidebar.go @@ -44,9 +44,10 @@ func (m *UI) modelInfo(width int) string { var modelContext *common.ModelContextInfo if model != nil && m.session != nil { modelContext = &common.ModelContextInfo{ - ContextUsed: m.session.CompletionTokens + m.session.PromptTokens, - Cost: m.session.Cost, - ModelContext: model.CatwalkCfg.ContextWindow, + ContextUsed: m.session.CompletionTokens + m.session.PromptTokens, + Cost: m.session.Cost, + ModelContext: model.CatwalkCfg.ContextWindow, + EstimatedUsage: m.session.EstimatedUsage, } } var modelName string diff --git a/internal/ui/styles/quickstyle.go b/internal/ui/styles/quickstyle.go index 7e01299788620df6c7bc9e25f913e8968c8fd110..cf685cf66f65a1abc388337a47e35ecbccc61c3d 100644 --- a/internal/ui/styles/quickstyle.go +++ b/internal/ui/styles/quickstyle.go @@ -763,6 +763,7 @@ func quickStyle(o quickStyleOpts) Styles { s.ModelInfo.Reasoning = lipgloss.NewStyle().Foreground(o.fgMostSubtle).PaddingLeft(2) s.ModelInfo.TokenCount = lipgloss.NewStyle().Foreground(o.fgMostSubtle) s.ModelInfo.TokenPercentage = lipgloss.NewStyle().Foreground(o.fgMoreSubtle) + s.ModelInfo.EstimatedUsagePrefix = lipgloss.NewStyle().Foreground(o.fgBase) s.ModelInfo.Cost = lipgloss.NewStyle().Foreground(o.fgMoreSubtle) s.ModelInfo.HypercreditIcon = lipgloss.NewStyle().Foreground(charmtone.Dolly) s.ModelInfo.HypercreditText = lipgloss.NewStyle().Foreground(o.fgMoreSubtle) diff --git a/internal/ui/styles/styles.go b/internal/ui/styles/styles.go index b8841e5c2881d25006840eb61d339bd104024901..20bb5d2424e774858045cdbf1bae836a518cde9e 100644 --- a/internal/ui/styles/styles.go +++ b/internal/ui/styles/styles.go @@ -183,16 +183,17 @@ type Styles struct { // ModelInfo (model name, provider, reasoning, token/cost summary) ModelInfo struct { - Icon lipgloss.Style // Model icon (◇) - Name lipgloss.Style // Model name text - Provider lipgloss.Style // "via " text - ProviderFallback lipgloss.Style // Provider on its own second line - Reasoning lipgloss.Style // Reasoning effort text - TokenCount lipgloss.Style // "(42K)" token count - TokenPercentage lipgloss.Style // "42%" percent of context window - Cost lipgloss.Style // "$0.42" cost readout - HypercreditIcon lipgloss.Style // Hypercredit icon (◆) - HypercreditText lipgloss.Style // Remaining Hypercredits text + Icon lipgloss.Style // Model icon (◇) + Name lipgloss.Style // Model name text + Provider lipgloss.Style // "via " text + ProviderFallback lipgloss.Style // Provider on its own second line + Reasoning lipgloss.Style // Reasoning effort text + TokenCount lipgloss.Style // "(42K)" token count + TokenPercentage lipgloss.Style // "42%" percent of context window + EstimatedUsagePrefix lipgloss.Style // "~" prefix for estimated usage + Cost lipgloss.Style // "$0.42" cost readout + HypercreditIcon lipgloss.Style // Hypercredit icon (◆) + HypercreditText lipgloss.Style // Remaining Hypercredits text } // Resource styles the LSP/MCP/skills sidebar lists: their heading,