Skip to content

Commit ed80256

Browse files
AnnatarHeAnnatarHe
authored
Merge pull request #190 from shelltime/feat/cc-statusline-command
feat(cc): add statusline command for Claude Code integration
2 parents 0737e79 + fb6be48 commit ed80256

12 files changed

Lines changed: 1907 additions & 5 deletions

‎commands/cc.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ var CCCommand = &cli.Command{
1212
Subcommands: []*cli.Command{
1313
CCInstallCommand,
1414
CCUninstallCommand,
15+
CCStatuslineCommand,
1516
},
1617
}
1718

‎commands/cc_statusline.go‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
package commands
2+
3+
import (
4+
"bufio"
5+
"context"
6+
"encoding/json"
7+
"fmt"
8+
"io"
9+
"os"
10+
"strings"
11+
"time"
12+
13+
"github.com/gookit/color"
14+
"github.com/malamtime/cli/daemon"
15+
"github.com/malamtime/cli/model"
16+
"github.com/urfave/cli/v2"
17+
)
18+
19+
var CCStatuslineCommand = &cli.Command{
20+
Name: "statusline",
21+
Usage: "Output statusline for Claude Code (reads JSON from stdin)",
22+
Action: commandCCStatusline,
23+
}
24+
25+
func commandCCStatusline(c *cli.Context) error {
26+
// Hard timeout for entire operation - statusline must be fast
27+
ctx, cancel := context.WithTimeout(c.Context, 100*time.Millisecond)
28+
defer cancel()
29+
30+
// Read from stdin
31+
input, err := readStdinWithTimeout(ctx)
32+
if err != nil {
33+
outputFallback()
34+
return nil
35+
}
36+
37+
// Parse input
38+
var data model.CCStatuslineInput
39+
if err := json.Unmarshal(input, &data); err != nil {
40+
outputFallback()
41+
return nil
42+
}
43+
44+
// Calculate context percentage
45+
contextPercent := calculateContextPercent(data.ContextWindow)
46+
47+
// Get daily cost - try daemon first, fallback to direct API
48+
var dailyCost float64
49+
config, err := configService.ReadConfigFile(ctx)
50+
if err == nil {
51+
dailyCost = getDailyCostWithDaemonFallback(ctx, config)
52+
}
53+
54+
// Format and output
55+
output := formatStatuslineOutput(data.Model.DisplayName, data.Cost.TotalCostUSD, dailyCost, contextPercent)
56+
fmt.Println(output)
57+
58+
return nil
59+
}
60+
61+
func readStdinWithTimeout(ctx context.Context) ([]byte, error) {
62+
resultCh := make(chan []byte, 1)
63+
errCh := make(chan error, 1)
64+
65+
go func() {
66+
reader := bufio.NewReader(os.Stdin)
67+
var data []byte
68+
for {
69+
line, err := reader.ReadBytes('\n')
70+
data = append(data, line...)
71+
if err != nil {
72+
if err == io.EOF {
73+
break
74+
}
75+
errCh <- err
76+
return
77+
}
78+
}
79+
resultCh <- data
80+
}()
81+
82+
select {
83+
case <-ctx.Done():
84+
return nil, ctx.Err()
85+
case err := <-errCh:
86+
return nil, err
87+
case data := <-resultCh:
88+
return data, nil
89+
}
90+
}
91+
92+
func calculateContextPercent(cw model.CCStatuslineContextWindow) float64 {
93+
if cw.ContextWindowSize == 0 {
94+
return 0
95+
}
96+
97+
// Use current_usage if available for accurate context window state
98+
if cw.CurrentUsage != nil {
99+
currentTokens := cw.CurrentUsage.InputTokens +
100+
cw.CurrentUsage.OutputTokens +
101+
cw.CurrentUsage.CacheCreationInputTokens +
102+
cw.CurrentUsage.CacheReadInputTokens
103+
return float64(currentTokens) / float64(cw.ContextWindowSize) * 100
104+
}
105+
106+
// Fallback to total tokens
107+
currentTokens := cw.TotalInputTokens + cw.TotalOutputTokens
108+
return float64(currentTokens) / float64(cw.ContextWindowSize) * 100
109+
}
110+
111+
func formatStatuslineOutput(modelName string, sessionCost, dailyCost, contextPercent float64) string {
112+
var parts []string
113+
114+
// Model name
115+
modelStr := fmt.Sprintf("🤖 %s", modelName)
116+
parts = append(parts, modelStr)
117+
118+
// Session cost (cyan)
119+
sessionStr := color.Cyan.Sprintf("💰 $%.2f", sessionCost)
120+
parts = append(parts, sessionStr)
121+
122+
// Daily cost (yellow)
123+
if dailyCost > 0 {
124+
dailyStr := color.Yellow.Sprintf("📊 $%.2f", dailyCost)
125+
parts = append(parts, dailyStr)
126+
} else {
127+
parts = append(parts, color.Gray.Sprint("📊 -"))
128+
}
129+
130+
// Context percentage with color coding
131+
var contextStr string
132+
switch {
133+
case contextPercent >= 80:
134+
contextStr = color.Red.Sprintf("📈 %.0f%%", contextPercent)
135+
case contextPercent >= 50:
136+
contextStr = color.Yellow.Sprintf("📈 %.0f%%", contextPercent)
137+
default:
138+
contextStr = color.Green.Sprintf("📈 %.0f%%", contextPercent)
139+
}
140+
parts = append(parts, contextStr)
141+
142+
return strings.Join(parts, " | ")
143+
}
144+
145+
func outputFallback() {
146+
fmt.Println(color.Gray.Sprint("🤖 - | 💰 - | 📊 - | 📈 -%"))
147+
}
148+
149+
// getDailyCostWithDaemonFallback tries to get daily cost from daemon first,
150+
// falls back to direct API if daemon is unavailable
151+
func getDailyCostWithDaemonFallback(ctx context.Context, config model.ShellTimeConfig) float64 {
152+
socketPath := config.SocketPath
153+
if socketPath == "" {
154+
socketPath = model.DefaultSocketPath
155+
}
156+
157+
// Try daemon first (50ms timeout for fast path)
158+
if daemon.IsSocketReady(ctx, socketPath) {
159+
resp, err := daemon.RequestCCInfo(socketPath, daemon.CCInfoTimeRangeToday, 50*time.Millisecond)
160+
if err == nil && resp != nil {
161+
return resp.TotalCostUSD
162+
}
163+
}
164+
165+
// Fallback to direct API (existing behavior)
166+
return model.FetchDailyCostCached(ctx, config)
167+
}

‎commands/cc_statusline_test.go‎

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
1+
package commands
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"net"
7+
"os"
8+
"path/filepath"
9+
"testing"
10+
"time"
11+
12+
"github.com/malamtime/cli/daemon"
13+
"github.com/malamtime/cli/model"
14+
"github.com/stretchr/testify/assert"
15+
"github.com/stretchr/testify/suite"
16+
)
17+
18+
type CCStatuslineTestSuite struct {
19+
suite.Suite
20+
mockConfig *model.MockConfigService
21+
origConfig model.ConfigService
22+
socketPath string
23+
listener net.Listener
24+
}
25+
26+
func (s *CCStatuslineTestSuite) SetupTest() {
27+
s.origConfig = configService
28+
s.mockConfig = model.NewMockConfigService(s.T())
29+
configService = s.mockConfig
30+
31+
s.socketPath = filepath.Join(os.TempDir(), "test-statusline.sock")
32+
os.Remove(s.socketPath)
33+
}
34+
35+
func (s *CCStatuslineTestSuite) TearDownTest() {
36+
configService = s.origConfig
37+
if s.listener != nil {
38+
s.listener.Close()
39+
}
40+
os.Remove(s.socketPath)
41+
}
42+
43+
// getDailyCostWithDaemonFallback Tests
44+
45+
func (s *CCStatuslineTestSuite) TestGetDailyCost_UsesDaemonWhenAvailable() {
46+
// Start mock daemon socket
47+
listener, err := net.Listen("unix", s.socketPath)
48+
assert.NoError(s.T(), err)
49+
s.listener = listener
50+
51+
expectedCost := 15.67
52+
go func() {
53+
conn, _ := listener.Accept()
54+
defer conn.Close()
55+
56+
var msg daemon.SocketMessage
57+
json.NewDecoder(conn).Decode(&msg)
58+
59+
response := daemon.CCInfoResponse{
60+
TotalCostUSD: expectedCost,
61+
TimeRange: "today",
62+
CachedAt: time.Now(),
63+
}
64+
json.NewEncoder(conn).Encode(response)
65+
}()
66+
67+
time.Sleep(10 * time.Millisecond)
68+
69+
config := model.ShellTimeConfig{
70+
SocketPath: s.socketPath,
71+
}
72+
73+
cost := getDailyCostWithDaemonFallback(context.Background(), config)
74+
75+
assert.Equal(s.T(), expectedCost, cost)
76+
}
77+
78+
func (s *CCStatuslineTestSuite) TestGetDailyCost_FallbackWhenDaemonUnavailable() {
79+
// No socket exists, should fall back to cached API
80+
config := model.ShellTimeConfig{
81+
SocketPath: "/nonexistent/socket.sock",
82+
Token: "", // No token means FetchDailyCostCached returns 0
83+
}
84+
85+
cost := getDailyCostWithDaemonFallback(context.Background(), config)
86+
87+
// Should return 0 (from cache fallback with no token)
88+
assert.Equal(s.T(), float64(0), cost)
89+
}
90+
91+
func (s *CCStatuslineTestSuite) TestGetDailyCost_FallbackOnDaemonError() {
92+
// Start mock daemon that returns error
93+
listener, err := net.Listen("unix", s.socketPath)
94+
assert.NoError(s.T(), err)
95+
s.listener = listener
96+
97+
go func() {
98+
conn, _ := listener.Accept()
99+
// Close immediately to cause error
100+
conn.Close()
101+
}()
102+
103+
time.Sleep(10 * time.Millisecond)
104+
105+
config := model.ShellTimeConfig{
106+
SocketPath: s.socketPath,
107+
Token: "", // No token
108+
}
109+
110+
cost := getDailyCostWithDaemonFallback(context.Background(), config)
111+
112+
// Should fall back and return 0
113+
assert.Equal(s.T(), float64(0), cost)
114+
}
115+
116+
func (s *CCStatuslineTestSuite) TestGetDailyCost_UsesDefaultSocketPath() {
117+
// Test that default socket path is used when config is empty
118+
config := model.ShellTimeConfig{
119+
SocketPath: "", // Empty path
120+
Token: "",
121+
}
122+
123+
// This should use model.DefaultSocketPath internally
124+
// Since no daemon is running, it will fall back
125+
cost := getDailyCostWithDaemonFallback(context.Background(), config)
126+
127+
assert.Equal(s.T(), float64(0), cost)
128+
}
129+
130+
// formatStatuslineOutput Tests
131+
132+
func (s *CCStatuslineTestSuite) TestFormatStatuslineOutput_AllValues() {
133+
output := formatStatuslineOutput("claude-opus-4", 1.23, 4.56, 75.0)
134+
135+
// Should contain all components
136+
assert.Contains(s.T(), output, "🤖 claude-opus-4")
137+
assert.Contains(s.T(), output, "$1.23")
138+
assert.Contains(s.T(), output, "$4.56")
139+
assert.Contains(s.T(), output, "75%") // Context percentage
140+
}
141+
142+
func (s *CCStatuslineTestSuite) TestFormatStatuslineOutput_ZeroDailyCost() {
143+
output := formatStatuslineOutput("claude-sonnet", 0.50, 0, 50.0)
144+
145+
// Should show "-" for zero daily cost
146+
assert.Contains(s.T(), output, "📊 -")
147+
}
148+
149+
func (s *CCStatuslineTestSuite) TestFormatStatuslineOutput_HighContextPercentage() {
150+
output := formatStatuslineOutput("test-model", 1.0, 1.0, 85.0)
151+
152+
// Should contain the percentage (color codes may vary)
153+
assert.Contains(s.T(), output, "85%")
154+
}
155+
156+
func (s *CCStatuslineTestSuite) TestFormatStatuslineOutput_LowContextPercentage() {
157+
output := formatStatuslineOutput("test-model", 1.0, 1.0, 25.0)
158+
159+
// Should contain the percentage
160+
assert.Contains(s.T(), output, "25%")
161+
}
162+
163+
// calculateContextPercent Tests
164+
165+
func (s *CCStatuslineTestSuite) TestCalculateContextPercent_ZeroContextWindowSize() {
166+
cw := model.CCStatuslineContextWindow{
167+
ContextWindowSize: 0,
168+
TotalInputTokens: 1000,
169+
TotalOutputTokens: 500,
170+
}
171+
172+
percent := calculateContextPercent(cw)
173+
174+
assert.Equal(s.T(), float64(0), percent)
175+
}
176+
177+
func (s *CCStatuslineTestSuite) TestCalculateContextPercent_WithCurrentUsage() {
178+
cw := model.CCStatuslineContextWindow{
179+
ContextWindowSize: 100000,
180+
CurrentUsage: &model.CCStatuslineContextUsage{
181+
InputTokens: 10000,
182+
OutputTokens: 5000,
183+
CacheCreationInputTokens: 2000,
184+
CacheReadInputTokens: 3000,
185+
},
186+
}
187+
188+
percent := calculateContextPercent(cw)
189+
190+
// (10000 + 5000 + 2000 + 3000) / 100000 * 100 = 20%
191+
assert.Equal(s.T(), float64(20), percent)
192+
}
193+
194+
func (s *CCStatuslineTestSuite) TestCalculateContextPercent_WithoutCurrentUsage() {
195+
cw := model.CCStatuslineContextWindow{
196+
ContextWindowSize: 100000,
197+
TotalInputTokens: 30000,
198+
TotalOutputTokens: 20000,
199+
CurrentUsage: nil,
200+
}
201+
202+
percent := calculateContextPercent(cw)
203+
204+
// (30000 + 20000) / 100000 * 100 = 50%
205+
assert.Equal(s.T(), float64(50), percent)
206+
}
207+
208+
func TestCCStatuslineTestSuite(t *testing.T) {
209+
suite.Run(t, new(CCStatuslineTestSuite))
210+
}

0 commit comments

Comments
 (0)