Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions commands/cc_statusline.go
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ func commandCCStatusline(c *cli.Context) error {
}
}

result = getDaemonInfoWithFallback(ctx, config, data.Cwd)
result = getDaemonInfoWithFallback(ctx, config, data.Cwd, data.Version)
}

// Format and output
Expand Down Expand Up @@ -304,15 +304,15 @@ func formatSessionDuration(totalSeconds int) string {

// getDaemonInfoWithFallback tries to get daily stats and git info from daemon first,
// falls back to direct API for stats if daemon is unavailable (git info only from daemon)
func getDaemonInfoWithFallback(ctx context.Context, config model.ShellTimeConfig, workingDir string) ccStatuslineResult {
func getDaemonInfoWithFallback(ctx context.Context, config model.ShellTimeConfig, workingDir, claudeCodeVersion string) ccStatuslineResult {
socketPath := config.SocketPath
if socketPath == "" {
socketPath = model.DefaultSocketPath
}

// Try daemon first (50ms timeout for fast path)
if daemon.IsSocketReady(ctx, socketPath) {
resp, err := daemon.RequestCCInfo(socketPath, daemon.CCInfoTimeRangeToday, workingDir, 50*time.Millisecond)
resp, err := daemon.RequestCCInfo(socketPath, daemon.CCInfoTimeRangeToday, workingDir, claudeCodeVersion, 50*time.Millisecond)
if err == nil && resp != nil {
return ccStatuslineResult{
Cost: resp.TotalCostUSD,
Expand Down
10 changes: 5 additions & 5 deletions commands/cc_statusline_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ func (s *CCStatuslineTestSuite) TestGetDaemonInfo_UsesDaemonWhenAvailable() {
SocketPath: s.socketPath,
}

result := getDaemonInfoWithFallback(context.Background(), config, "/some/path")
result := getDaemonInfoWithFallback(context.Background(), config, "/some/path", "")

assert.Equal(s.T(), expectedCost, result.Cost)
assert.Equal(s.T(), expectedSessionSeconds, result.SessionSeconds)
Expand All @@ -94,7 +94,7 @@ func (s *CCStatuslineTestSuite) TestGetDaemonInfo_FallbackWhenDaemonUnavailable(
Token: "", // No token means FetchDailyStatsCached returns zero values
}

result := getDaemonInfoWithFallback(context.Background(), config, "")
result := getDaemonInfoWithFallback(context.Background(), config, "", "")

// Should return zero values (from cache fallback with no token)
assert.Equal(s.T(), float64(0), result.Cost)
Expand Down Expand Up @@ -122,7 +122,7 @@ func (s *CCStatuslineTestSuite) TestGetDaemonInfo_FallbackOnDaemonError() {
Token: "", // No token
}

result := getDaemonInfoWithFallback(context.Background(), config, "")
result := getDaemonInfoWithFallback(context.Background(), config, "", "")

// Should fall back and return zero values
assert.Equal(s.T(), float64(0), result.Cost)
Expand All @@ -141,7 +141,7 @@ func (s *CCStatuslineTestSuite) TestGetDaemonInfo_UsesDefaultSocketPath() {
// This should use model.DefaultSocketPath internally
// Since no daemon is running at the default path, it will fall back to cached API
// The function should not panic and should return a valid result struct
result := getDaemonInfoWithFallback(context.Background(), config, "")
result := getDaemonInfoWithFallback(context.Background(), config, "", "")

// We can't assert on exact values since the global cache might have data
// from previous tests. Just verify the function returns without error
Expand Down Expand Up @@ -582,7 +582,7 @@ func (s *CCStatuslineTestSuite) TestGetDaemonInfo_PropagatesRateLimitFields() {
SocketPath: s.socketPath,
}

result := getDaemonInfoWithFallback(context.Background(), config, "/some/path")
result := getDaemonInfoWithFallback(context.Background(), config, "/some/path", "")

assert.NotNil(s.T(), result.FiveHourUtilization)
assert.NotNil(s.T(), result.SevenDayUtilization)
Expand Down
59 changes: 55 additions & 4 deletions daemon/anthropic_ratelimit.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,26 @@ import (
"os/exec"
"path/filepath"
"runtime"
"strconv"
"strings"
"sync"
"time"
)

const anthropicUsageCacheTTL = 10 * time.Minute

// anthropicRateLimitBackoff is the minimum cooldown applied after a 429 from the usage API,
// used when the response carries no (or a shorter) Retry-After. It is longer than the normal
// TTL so the daemon stops poking the throttled bucket.
const anthropicRateLimitBackoff = 30 * time.Minute

// claudeCodeFallbackVersion is used in the User-Agent when the real Claude Code version is
// unknown. The usage endpoint gates on the "claude-code/" prefix, not the exact version.
const claudeCodeFallbackVersion = "2.0.0"

// anthropicUsageURL is the OAuth usage endpoint. It is a var (not a const) so tests can override it.
var anthropicUsageURL = "https://api.anthropic.com/api/oauth/usage"

// AnthropicRateLimitData holds the parsed rate limit utilization data
type AnthropicRateLimitData struct {
FiveHourUtilization float64
Expand All @@ -28,7 +42,19 @@ type anthropicRateLimitCache struct {
usage *AnthropicRateLimitData
fetchedAt time.Time
lastAttemptAt time.Time
lastError string // short error description for statusline display
backoffUntil time.Time // when set in the future, skip fetching (e.g. after a 429)
lastError string // short error description for statusline display
}

// anthropicAPIError represents a non-200 response from the Anthropic usage API.
// Error() keeps the historical message format so shortenAPIError still produces "api:<code>".
type anthropicAPIError struct {
StatusCode int
RetryAfter time.Duration // parsed from the Retry-After header; 0 if absent/unparseable
}

func (e *anthropicAPIError) Error() string {
return fmt.Sprintf("anthropic usage API returned status %d", e.StatusCode)
}

// anthropicUsageResponse maps the Anthropic API response
Expand Down Expand Up @@ -111,14 +137,23 @@ func parseOAuthTokenFromJSON(data []byte) (string, error) {
}

// fetchAnthropicUsage calls the Anthropic OAuth usage API and returns rate limit data.
func fetchAnthropicUsage(ctx context.Context, token string) (*AnthropicRateLimitData, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://api.anthropic.com/api/oauth/usage", nil)
// version is the Claude Code version used for the User-Agent header; when empty it falls back
// to claudeCodeFallbackVersion. The endpoint requires a "claude-code/<version>" User-Agent:
// without it, requests land in an aggressively rate-limited bucket and return persistent 429s.
func fetchAnthropicUsage(ctx context.Context, token, version string) (*AnthropicRateLimitData, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, anthropicUsageURL, nil)
if err != nil {
return nil, err
}

if version == "" {
version = claudeCodeFallbackVersion
}

req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("anthropic-beta", "oauth-2025-04-20")
req.Header.Set("User-Agent", "claude-code/"+version)
req.Header.Set("Content-Type", "application/json")

client := &http.Client{Timeout: 5 * time.Second}
resp, err := client.Do(req)
Expand All @@ -128,7 +163,10 @@ func fetchAnthropicUsage(ctx context.Context, token string) (*AnthropicRateLimit
defer resp.Body.Close()

if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("anthropic usage API returned status %d", resp.StatusCode)
return nil, &anthropicAPIError{
StatusCode: resp.StatusCode,
RetryAfter: parseRetryAfter(resp.Header.Get("Retry-After")),
}
}

var usage anthropicUsageResponse
Expand All @@ -143,3 +181,16 @@ func fetchAnthropicUsage(ctx context.Context, token string) (*AnthropicRateLimit
SevenDayResetsAt: usage.SevenDay.ResetsAt,
}, nil
}

// parseRetryAfter parses an HTTP Retry-After header in delay-seconds form.
// Returns 0 when the value is absent or not a positive integer (HTTP-date form is not used here).
func parseRetryAfter(v string) time.Duration {
v = strings.TrimSpace(v)
if v == "" {
return 0
}
if secs, err := strconv.Atoi(v); err == nil && secs > 0 {
return time.Duration(secs) * time.Second
}
return 0
}
Comment on lines +187 to +196

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To prevent potential integer overflow and excessively long backoff periods, it is highly recommended to cap the parsed Retry-After value to a reasonable maximum duration (e.g., 24 hours). If a buggy or malicious server returns an extremely large value, it could cause time.Duration overflow or result in the daemon backing off indefinitely.

Suggested change
func parseRetryAfter(v string) time.Duration {
v = strings.TrimSpace(v)
if v == "" {
return 0
}
if secs, err := strconv.Atoi(v); err == nil && secs > 0 {
return time.Duration(secs) * time.Second
}
return 0
}
func parseRetryAfter(v string) time.Duration {
v = strings.TrimSpace(v)
if v == "" {
return 0
}
if secs, err := strconv.Atoi(v); err == nil && secs > 0 {
const maxSeconds = 24 * 3600
if secs > maxSeconds {
secs = maxSeconds
}
return time.Duration(secs) * time.Second
}
return 0
}

120 changes: 120 additions & 0 deletions daemon/anthropic_ratelimit_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,137 @@ package daemon
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"testing"
"time"

"github.com/malamtime/cli/model"
"github.com/stretchr/testify/assert"
)

// withTestUsageURL points fetchAnthropicUsage at a test server for the duration of the test.
func withTestUsageURL(t *testing.T, url string) {
t.Helper()
orig := anthropicUsageURL
anthropicUsageURL = url
t.Cleanup(func() { anthropicUsageURL = orig })
}

func TestFetchAnthropicUsage_SetsClaudeCodeUserAgent(t *testing.T) {
var gotUA, gotCT string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotUA = r.Header.Get("User-Agent")
gotCT = r.Header.Get("Content-Type")
assert.Equal(t, "Bearer tok", r.Header.Get("Authorization"))
assert.Equal(t, "oauth-2025-04-20", r.Header.Get("anthropic-beta"))
json.NewEncoder(w).Encode(anthropicUsageResponse{})
}))
defer server.Close()
withTestUsageURL(t, server.URL)

_, err := fetchAnthropicUsage(context.Background(), "tok", "9.9.9")
assert.NoError(t, err)
assert.Equal(t, "claude-code/9.9.9", gotUA)
assert.Equal(t, "application/json", gotCT)
}

func TestFetchAnthropicUsage_UserAgentFallback(t *testing.T) {
var gotUA string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotUA = r.Header.Get("User-Agent")
json.NewEncoder(w).Encode(anthropicUsageResponse{})
}))
defer server.Close()
withTestUsageURL(t, server.URL)

_, err := fetchAnthropicUsage(context.Background(), "tok", "")
assert.NoError(t, err)
assert.Equal(t, "claude-code/"+claudeCodeFallbackVersion, gotUA)
}

func TestFetchAnthropicUsage_429ReturnsTypedErrorWithRetryAfter(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Retry-After", "120")
w.WriteHeader(http.StatusTooManyRequests)
}))
defer server.Close()
withTestUsageURL(t, server.URL)

_, err := fetchAnthropicUsage(context.Background(), "tok", "1.0.0")
assert.Error(t, err)

var apiErr *anthropicAPIError
assert.True(t, errors.As(err, &apiErr))
assert.Equal(t, http.StatusTooManyRequests, apiErr.StatusCode)
assert.Equal(t, 120*time.Second, apiErr.RetryAfter)

// The typed error must still shorten to the historical "api:429" for the statusline.
assert.Equal(t, "api:429", shortenAPIError(err))
}

func TestParseRetryAfter(t *testing.T) {
cases := []struct {
in string
want time.Duration
}{
{"", 0},
{"120", 120 * time.Second},
{" 60 ", 60 * time.Second},
{"0", 0},
{"-5", 0},
{"abc", 0},
{"Wed, 21 Oct 2026 07:28:00 GMT", 0}, // HTTP-date form is not parsed
}
for _, c := range cases {
assert.Equalf(t, c.want, parseRetryAfter(c.in), "parseRetryAfter(%q)", c.in)
}
}

func TestSetGetClaudeCodeVersion(t *testing.T) {
service := NewCCInfoTimerService(&model.ShellTimeConfig{})
assert.Empty(t, service.GetClaudeCodeVersion())

service.SetClaudeCodeVersion("3.1.4")
assert.Equal(t, "3.1.4", service.GetClaudeCodeVersion())

// Empty values are ignored so a known version is not clobbered.
service.SetClaudeCodeVersion("")
assert.Equal(t, "3.1.4", service.GetClaudeCodeVersion())
}

func TestStopTimer_PreservesRateLimitCache(t *testing.T) {
service := NewCCInfoTimerService(&model.ShellTimeConfig{})

// Seed a good rate-limit cache.
service.rateLimitCache.mu.Lock()
service.rateLimitCache.usage = &AnthropicRateLimitData{FiveHourUtilization: 0.5}
service.rateLimitCache.fetchedAt = time.Now()
service.rateLimitCache.lastAttemptAt = time.Now()
service.rateLimitCache.backoffUntil = time.Now().Add(time.Hour)
service.rateLimitCache.mu.Unlock()

// stopTimer must be called with timerMu held and the timer "running".
service.timerMu.Lock()
service.timerRunning = true
service.ticker = time.NewTicker(time.Hour)
service.stopTimer()
service.timerMu.Unlock()

// The cache must survive an idle stop so the TTL/backoff hold across idle cycles.
rl := service.GetCachedRateLimit()
assert.NotNil(t, rl)
assert.Equal(t, 0.5, rl.FiveHourUtilization)

service.rateLimitCache.mu.RLock()
assert.False(t, service.rateLimitCache.backoffUntil.IsZero())
service.rateLimitCache.mu.RUnlock()
}

func TestFetchAnthropicUsage_Success(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "GET", r.Method)
Expand Down
10 changes: 5 additions & 5 deletions daemon/cc_info_handler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -232,7 +232,7 @@ func (s *CCInfoClientTestSuite) TestRequestCCInfo_Success() {
// Give server time to start
time.Sleep(10 * time.Millisecond)

response, err := RequestCCInfo(s.socketPath, CCInfoTimeRangeToday, "", 1*time.Second)
response, err := RequestCCInfo(s.socketPath, CCInfoTimeRangeToday, "", "", 1*time.Second)

assert.NoError(s.T(), err)
assert.NotNil(s.T(), response)
Expand All @@ -255,14 +255,14 @@ func (s *CCInfoClientTestSuite) TestRequestCCInfo_Timeout() {

time.Sleep(10 * time.Millisecond)

response, err := RequestCCInfo(s.socketPath, CCInfoTimeRangeToday, "", 50*time.Millisecond)
response, err := RequestCCInfo(s.socketPath, CCInfoTimeRangeToday, "", "", 50*time.Millisecond)

assert.Error(s.T(), err)
assert.Nil(s.T(), response)
}

func (s *CCInfoClientTestSuite) TestRequestCCInfo_SocketNotFound() {
response, err := RequestCCInfo("/nonexistent/socket.sock", CCInfoTimeRangeToday, "", 100*time.Millisecond)
response, err := RequestCCInfo("/nonexistent/socket.sock", CCInfoTimeRangeToday, "", "", 100*time.Millisecond)

assert.Error(s.T(), err)
assert.Nil(s.T(), response)
Expand All @@ -287,7 +287,7 @@ func (s *CCInfoClientTestSuite) TestRequestCCInfo_InvalidResponse() {

time.Sleep(10 * time.Millisecond)

response, err := RequestCCInfo(s.socketPath, CCInfoTimeRangeToday, "", 1*time.Second)
response, err := RequestCCInfo(s.socketPath, CCInfoTimeRangeToday, "", "", 1*time.Second)

assert.Error(s.T(), err)
assert.Nil(s.T(), response)
Expand All @@ -312,7 +312,7 @@ func (s *CCInfoClientTestSuite) TestRequestCCInfo_SendsCorrectMessage() {

time.Sleep(10 * time.Millisecond)

RequestCCInfo(s.socketPath, CCInfoTimeRangeWeek, "", 1*time.Second)
RequestCCInfo(s.socketPath, CCInfoTimeRangeWeek, "", "", 1*time.Second)

assert.Equal(s.T(), SocketMessageTypeCCInfo, receivedMsg.Type)

Expand Down
Loading
Loading