mirror of
https://github.com/JanDeDobbeleer/oh-my-posh.git
synced 2026-08-24 10:14:12 -05:00
fix(claude): reset gauge to 0 when context is cleared
This commit is contained in:
committed by
Jan De Dobbeleer
parent
8462c168bc
commit
072b71884d
+64
-20
@@ -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
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user