fix(claude): reset gauge to 0 when context is cleared

This commit is contained in:
Scott Beardsley
2026-02-06 07:42:20 +01:00
committed by Jan De Dobbeleer
parent 8462c168bc
commit 072b71884d
2 changed files with 236 additions and 49 deletions
+64 -20
View File
@@ -16,11 +16,11 @@ type Claude struct {
// ClaudeData represents the parsed Claude JSON data
type ClaudeData struct {
SessionID string `json:"session_id"`
Model ClaudeModel `json:"model"`
Workspace ClaudeWorkspace `json:"workspace"`
Cost ClaudeCost `json:"cost"`
SessionID string `json:"session_id"`
ContextWindow ClaudeContextWindow `json:"context_window"`
Cost ClaudeCost `json:"cost"`
}
// ClaudeModel represents the AI model information
@@ -43,16 +43,20 @@ type ClaudeCost struct {
// ClaudeContextWindow represents token usage information
type ClaudeContextWindow struct {
TotalInputTokens int `json:"total_input_tokens"`
TotalOutputTokens int `json:"total_output_tokens"`
ContextWindowSize int `json:"context_window_size"`
CurrentUsage ClaudeCurrentUsage `json:"current_usage"`
UsedPercentage *int `json:"used_percentage"`
RemainingPercentage *int `json:"remaining_percentage"`
CurrentUsage *ClaudeCurrentUsage `json:"current_usage"`
TotalInputTokens int `json:"total_input_tokens"`
TotalOutputTokens int `json:"total_output_tokens"`
ContextWindowSize int `json:"context_window_size"`
}
// ClaudeCurrentUsage represents current message token usage
// ClaudeCurrentUsage represents current context window usage from the last API call
type ClaudeCurrentUsage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
CacheCreationInputTokens int `json:"cache_creation_input_tokens"`
CacheReadInputTokens int `json:"cache_read_input_tokens"`
}
const (
@@ -83,19 +87,44 @@ func (c *Claude) Enabled() bool {
return true
}
// TokenUsagePercent returns the percentage of context window used by total tokens
// TokenUsagePercent returns the percentage of context window used.
// Uses pre-calculated UsedPercentage when available (resets on compact/clear),
// falls back to calculating from CurrentUsage, then to total tokens for backwards compatibility.
func (c *Claude) TokenUsagePercent() text.Percentage {
// Prefer pre-calculated UsedPercentage - most accurate and resets on compact/clear
// When UsedPercentage is nil (null in JSON), context was reset - return 0
if c.ContextWindow.UsedPercentage != nil {
if *c.ContextWindow.UsedPercentage > 100 {
return 100
}
return text.Percentage(*c.ContextWindow.UsedPercentage)
}
// UsedPercentage is nil - check if CurrentUsage is also nil (indicates reset/clear)
if c.ContextWindow.CurrentUsage == nil {
return 0
}
if c.ContextWindow.ContextWindowSize <= 0 {
return 0
}
totalTokens := c.ContextWindow.TotalInputTokens + c.ContextWindow.TotalOutputTokens
if totalTokens <= 0 {
// Calculate from CurrentUsage (includes cache tokens for accurate context measurement)
currentTokens := c.ContextWindow.CurrentUsage.InputTokens +
c.ContextWindow.CurrentUsage.CacheCreationInputTokens +
c.ContextWindow.CurrentUsage.CacheReadInputTokens
// Fallback to total tokens if CurrentUsage is not provided (backwards compatibility)
if currentTokens <= 0 {
currentTokens = c.ContextWindow.TotalInputTokens + c.ContextWindow.TotalOutputTokens
}
if currentTokens <= 0 {
return 0
}
// Use floating-point arithmetic for accurate percentage calculation
percent := (float64(totalTokens) * 100.0) / float64(c.ContextWindow.ContextWindowSize)
percent := (float64(currentTokens) * 100.0) / float64(c.ContextWindow.ContextWindowSize)
// Round to nearest integer and cap at 100
roundedPercent := int(percent + 0.5)
@@ -115,17 +144,32 @@ func (c *Claude) FormattedCost() string {
return fmt.Sprintf("$%.2f", c.Cost.TotalCostUSD)
}
// FormattedTokens returns a human-readable string of total tokens used
// FormattedTokens returns a human-readable string of current context tokens.
// Uses CurrentUsage (which represents actual context and resets on compact/clear)
// with fallback to total tokens for backwards compatibility.
func (c *Claude) FormattedTokens() string {
totalTokens := c.ContextWindow.TotalInputTokens + c.ContextWindow.TotalOutputTokens
var currentTokens int
if totalTokens < int(thousand) {
return fmt.Sprintf("%d", totalTokens)
// Use CurrentUsage for display - includes cache tokens for accurate context measurement
// When CurrentUsage is nil (context reset), fall back to total tokens
if c.ContextWindow.CurrentUsage != nil {
currentTokens = c.ContextWindow.CurrentUsage.InputTokens +
c.ContextWindow.CurrentUsage.CacheCreationInputTokens +
c.ContextWindow.CurrentUsage.CacheReadInputTokens
}
if totalTokens < int(million) {
return fmt.Sprintf("%.1fK", float64(totalTokens)/thousand)
// Fallback to total tokens if CurrentUsage is not provided (backwards compatibility)
if currentTokens <= 0 {
currentTokens = c.ContextWindow.TotalInputTokens + c.ContextWindow.TotalOutputTokens
}
return fmt.Sprintf("%.1fM", float64(totalTokens)/million)
if currentTokens < int(thousand) {
return fmt.Sprintf("%d", currentTokens)
}
if currentTokens < int(million) {
return fmt.Sprintf("%.1fK", float64(currentTokens)/thousand)
}
return fmt.Sprintf("%.1fM", float64(currentTokens)/million)
}
+172 -29
View File
@@ -44,7 +44,7 @@ func TestClaudeSegment(t *testing.T) {
TotalInputTokens: 15234,
TotalOutputTokens: 4521,
ContextWindowSize: 200000,
CurrentUsage: ClaudeCurrentUsage{
CurrentUsage: &ClaudeCurrentUsage{
InputTokens: 8500,
OutputTokens: 1200,
},
@@ -102,28 +102,63 @@ func TestClaudeSegment(t *testing.T) {
func TestClaudeTokenUsagePercent(t *testing.T) {
cases := []struct {
Case string
InputTokens int
OutputTokens int
ContextWindow int
ExpectedPercent text.Percentage
UsedPercentage *int
Case string
InputTokens int
OutputTokens int
CurrentInput int
CacheCreationInputTokens int
CacheReadInputTokens int
ContextWindow int
ExpectedPercent text.Percentage
HasCurrentUsage bool
}{
{
Case: "Zero context window",
Case: "Uses UsedPercentage when available",
UsedPercentage: intPtr(42),
ContextWindow: 200000,
ExpectedPercent: 42,
},
{
Case: "UsedPercentage capped at 100",
UsedPercentage: intPtr(150),
ContextWindow: 200000,
ExpectedPercent: 100,
},
{
Case: "UsedPercentage zero is valid",
UsedPercentage: intPtr(0),
ContextWindow: 200000,
ExpectedPercent: 0,
},
{
Case: "Context reset - both UsedPercentage and CurrentUsage nil",
UsedPercentage: nil,
HasCurrentUsage: false,
InputTokens: 50000, // High cumulative total - should be ignored
OutputTokens: 50000,
ContextWindow: 200000,
ExpectedPercent: 0, // Should return 0 after reset, not fallback to total
},
{
Case: "Zero context window (no UsedPercentage)",
HasCurrentUsage: true,
InputTokens: 1000,
OutputTokens: 500,
ContextWindow: 0,
ExpectedPercent: 0,
},
{
Case: "10% usage",
Case: "10% usage (fallback to total)",
HasCurrentUsage: true,
InputTokens: 8000,
OutputTokens: 2000,
ContextWindow: 100000,
ExpectedPercent: 10,
},
{
Case: "50% usage",
Case: "50% usage (fallback to total)",
HasCurrentUsage: true,
InputTokens: 50000,
OutputTokens: 50000,
ContextWindow: 200000,
@@ -131,17 +166,64 @@ func TestClaudeTokenUsagePercent(t *testing.T) {
},
{
Case: "Over 100% usage (capped)",
HasCurrentUsage: true,
InputTokens: 100000,
OutputTokens: 50000,
ContextWindow: 100000,
ExpectedPercent: 100,
},
{
Case: "Uses CurrentUsage input tokens",
HasCurrentUsage: true,
InputTokens: 100000, // High cumulative total
OutputTokens: 50000,
CurrentInput: 20000, // Current context input
ContextWindow: 200000,
ExpectedPercent: 10, // Should use current input (20000/200000 = 10%)
},
{
Case: "Uses CurrentUsage with cache tokens",
HasCurrentUsage: true,
InputTokens: 100000, // High cumulative total
OutputTokens: 50000,
CurrentInput: 10000,
CacheCreationInputTokens: 5000,
CacheReadInputTokens: 5000,
ContextWindow: 200000,
ExpectedPercent: 10, // (10000+5000+5000)/200000 = 10%
},
{
Case: "Uses CurrentUsage after compact (low current, high total)",
HasCurrentUsage: true,
InputTokens: 100000, // High cumulative total
OutputTokens: 50000,
CurrentInput: 6000, // Low current context (after compact)
ContextWindow: 200000,
ExpectedPercent: 3, // Should use current (6000/200000 = 3%)
},
{
Case: "Fallback to total when CurrentUsage is zero",
HasCurrentUsage: true,
InputTokens: 20000,
OutputTokens: 10000,
CurrentInput: 0,
ContextWindow: 100000,
ExpectedPercent: 30, // Should fallback to total (30000/100000 = 30%)
},
}
for _, tc := range cases {
claude := &Claude{}
claude.ContextWindow.TotalInputTokens = tc.InputTokens
claude.ContextWindow.TotalOutputTokens = tc.OutputTokens
if tc.HasCurrentUsage {
claude.ContextWindow.CurrentUsage = &ClaudeCurrentUsage{
InputTokens: tc.CurrentInput,
CacheCreationInputTokens: tc.CacheCreationInputTokens,
CacheReadInputTokens: tc.CacheReadInputTokens,
}
}
claude.ContextWindow.UsedPercentage = tc.UsedPercentage
claude.ContextWindow.ContextWindowSize = tc.ContextWindow
percent := claude.TokenUsagePercent()
@@ -149,6 +231,11 @@ func TestClaudeTokenUsagePercent(t *testing.T) {
}
}
// intPtr is a helper to create a pointer to an int value
func intPtr(i int) *int {
return &i
}
func TestClaudeFormattedCost(t *testing.T) {
cases := []struct {
Case string
@@ -188,34 +275,83 @@ func TestClaudeFormattedCost(t *testing.T) {
func TestClaudeFormattedTokens(t *testing.T) {
cases := []struct {
Case string
ExpectedFormat string
InputTokens int
OutputTokens int
Case string
ExpectedFormat string
InputTokens int
OutputTokens int
CurrentInput int
CacheCreationInputTokens int
CacheReadInputTokens int
HasCurrentUsage bool
}{
{
Case: "Small token count",
InputTokens: 300,
OutputTokens: 200,
ExpectedFormat: "500",
Case: "Small token count (fallback to total)",
HasCurrentUsage: true,
InputTokens: 300,
OutputTokens: 200,
ExpectedFormat: "500",
},
{
Case: "Thousands",
InputTokens: 8500,
OutputTokens: 1500,
ExpectedFormat: "10.0K",
Case: "Thousands (fallback to total)",
HasCurrentUsage: true,
InputTokens: 8500,
OutputTokens: 1500,
ExpectedFormat: "10.0K",
},
{
Case: "Tens of thousands",
InputTokens: 50000,
OutputTokens: 25000,
ExpectedFormat: "75.0K",
Case: "Tens of thousands (fallback to total)",
HasCurrentUsage: true,
InputTokens: 50000,
OutputTokens: 25000,
ExpectedFormat: "75.0K",
},
{
Case: "Millions",
InputTokens: 1500000,
OutputTokens: 500000,
ExpectedFormat: "2.0M",
Case: "Millions (fallback to total)",
HasCurrentUsage: true,
InputTokens: 1500000,
OutputTokens: 500000,
ExpectedFormat: "2.0M",
},
{
Case: "Uses CurrentUsage input tokens",
HasCurrentUsage: true,
InputTokens: 100000, // High cumulative total
OutputTokens: 50000,
CurrentInput: 10000, // Current context input
ExpectedFormat: "10.0K",
},
{
Case: "Uses CurrentUsage with cache tokens",
HasCurrentUsage: true,
InputTokens: 100000, // High cumulative total
OutputTokens: 50000,
CurrentInput: 5000,
CacheCreationInputTokens: 2500,
CacheReadInputTokens: 2500,
ExpectedFormat: "10.0K", // 5000+2500+2500 = 10000
},
{
Case: "Uses CurrentUsage after compact (low current)",
HasCurrentUsage: true,
InputTokens: 500000, // High cumulative total
OutputTokens: 200000,
CurrentInput: 500, // Low current context (after compact)
ExpectedFormat: "500",
},
{
Case: "Fallback to total when CurrentUsage is zero",
HasCurrentUsage: true,
InputTokens: 50000,
OutputTokens: 25000,
CurrentInput: 0,
ExpectedFormat: "75.0K", // Should fallback to total
},
{
Case: "Nil CurrentUsage falls back to total",
HasCurrentUsage: false,
InputTokens: 50000,
OutputTokens: 25000,
ExpectedFormat: "75.0K", // Should fallback to total
},
}
@@ -223,6 +359,13 @@ func TestClaudeFormattedTokens(t *testing.T) {
claude := &Claude{}
claude.ContextWindow.TotalInputTokens = tc.InputTokens
claude.ContextWindow.TotalOutputTokens = tc.OutputTokens
if tc.HasCurrentUsage {
claude.ContextWindow.CurrentUsage = &ClaudeCurrentUsage{
InputTokens: tc.CurrentInput,
CacheCreationInputTokens: tc.CacheCreationInputTokens,
CacheReadInputTokens: tc.CacheReadInputTokens,
}
}
formatted := claude.FormattedTokens()
assert.Equal(t, tc.ExpectedFormat, formatted, tc.Case)