Skip to content

Commit 16df58b

Browse files
committed
feat(daemon): add Codex usage rate limit fetching and reporting
Add support for fetching Codex rate limit usage data from the OpenAI API (~/.codex/auth.json) alongside existing Anthropic usage tracking. The daemon now loads Codex auth credentials, fetches usage every 10 minutes (same cache TTL as Anthropic), and sends the data to the server via POST /api/v1/codex-usage. https://claude.ai/code/session_01FPxeHEA8w2cC136GFpi3to
1 parent c6dac1a commit 16df58b

2 files changed

Lines changed: 313 additions & 1 deletion

File tree

daemon/cc_info_timer.go

Lines changed: 144 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,9 @@ type CCInfoTimerService struct {
5555
// Anthropic rate limit cache
5656
rateLimitCache *anthropicRateLimitCache
5757

58+
// Codex rate limit cache
59+
codexRateLimitCache *codexRateLimitCache
60+
5861
// User profile cache (permanent for daemon lifetime)
5962
userLogin string
6063
userLoginFetched bool
@@ -67,7 +70,8 @@ func NewCCInfoTimerService(config *model.ShellTimeConfig) *CCInfoTimerService {
6770
cache: make(map[CCInfoTimeRange]CCInfoCache),
6871
activeRanges: make(map[CCInfoTimeRange]bool),
6972
gitCache: make(map[string]*GitCacheEntry),
70-
rateLimitCache: &anthropicRateLimitCache{},
73+
rateLimitCache: &anthropicRateLimitCache{},
74+
codexRateLimitCache: &codexRateLimitCache{},
7175
stopChan: make(chan struct{}),
7276
}
7377
}
@@ -152,6 +156,11 @@ func (s *CCInfoTimerService) stopTimer() {
152156
s.rateLimitCache.fetchedAt = time.Time{}
153157
s.rateLimitCache.lastAttemptAt = time.Time{}
154158
s.rateLimitCache.mu.Unlock()
159+
s.codexRateLimitCache.mu.Lock()
160+
s.codexRateLimitCache.usage = nil
161+
s.codexRateLimitCache.fetchedAt = time.Time{}
162+
s.codexRateLimitCache.lastAttemptAt = time.Time{}
163+
s.codexRateLimitCache.mu.Unlock()
155164

156165
slog.Info("CC info timer stopped due to inactivity")
157166
}
@@ -171,6 +180,7 @@ func (s *CCInfoTimerService) timerLoop() {
171180
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
172181
defer cancel()
173182
s.fetchRateLimit(ctx)
183+
s.fetchCodexRateLimit(ctx)
174184
}()
175185
go s.fetchUserProfile(context.Background())
176186

@@ -194,6 +204,7 @@ func (s *CCInfoTimerService) timerLoop() {
194204
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
195205
defer cancel()
196206
s.fetchRateLimit(ctx)
207+
s.fetchCodexRateLimit(ctx)
197208
}()
198209

199210
case <-s.stopChan:
@@ -551,6 +562,138 @@ func (s *CCInfoTimerService) GetCachedRateLimitError() string {
551562
return s.rateLimitCache.lastError
552563
}
553564

565+
// fetchCodexRateLimit fetches Codex rate limit data if cache is stale.
566+
func (s *CCInfoTimerService) fetchCodexRateLimit(ctx context.Context) {
567+
if runtime.GOOS != "darwin" && runtime.GOOS != "linux" {
568+
return
569+
}
570+
571+
// Check cache TTL under read lock
572+
s.codexRateLimitCache.mu.RLock()
573+
sinceLastFetch := time.Since(s.codexRateLimitCache.fetchedAt)
574+
sinceLastAttempt := time.Since(s.codexRateLimitCache.lastAttemptAt)
575+
s.codexRateLimitCache.mu.RUnlock()
576+
577+
if sinceLastFetch < codexUsageCacheTTL || sinceLastAttempt < codexUsageCacheTTL {
578+
return
579+
}
580+
581+
// Record attempt time
582+
s.codexRateLimitCache.mu.Lock()
583+
s.codexRateLimitCache.lastAttemptAt = time.Now()
584+
s.codexRateLimitCache.mu.Unlock()
585+
586+
auth, err := loadCodexAuth()
587+
if err != nil || auth == nil {
588+
slog.Debug("Failed to load Codex auth", slog.Any("err", err))
589+
s.codexRateLimitCache.mu.Lock()
590+
s.codexRateLimitCache.lastError = "auth"
591+
s.codexRateLimitCache.mu.Unlock()
592+
return
593+
}
594+
595+
usage, err := fetchCodexUsage(ctx, auth)
596+
if err != nil {
597+
slog.Warn("Failed to fetch Codex usage", slog.Any("err", err))
598+
s.codexRateLimitCache.mu.Lock()
599+
s.codexRateLimitCache.lastError = shortenCodexAPIError(err)
600+
s.codexRateLimitCache.mu.Unlock()
601+
return
602+
}
603+
604+
s.codexRateLimitCache.mu.Lock()
605+
s.codexRateLimitCache.usage = usage
606+
s.codexRateLimitCache.fetchedAt = time.Now()
607+
s.codexRateLimitCache.lastError = ""
608+
s.codexRateLimitCache.mu.Unlock()
609+
610+
// Send usage data to server (fire-and-forget)
611+
go func() {
612+
bgCtx, bgCancel := context.WithTimeout(context.Background(), 10*time.Second)
613+
defer bgCancel()
614+
s.sendCodexUsageToServer(bgCtx, usage)
615+
}()
616+
617+
slog.Debug("Codex rate limit updated",
618+
slog.String("plan", usage.Plan),
619+
slog.Int("windows", len(usage.Windows)))
620+
}
621+
622+
// sendCodexUsageToServer sends Codex usage data to the ShellTime server
623+
// for scheduling push notifications when rate limits reset.
624+
func (s *CCInfoTimerService) sendCodexUsageToServer(ctx context.Context, usage *CodexRateLimitData) {
625+
if s.config.Token == "" {
626+
return
627+
}
628+
629+
type usageWindow struct {
630+
LimitID string `json:"limit_id"`
631+
UsagePercentage float64 `json:"usage_percentage"`
632+
ResetsAt string `json:"resets_at"`
633+
WindowDurationMinutes int `json:"window_duration_minutes"`
634+
}
635+
type usagePayload struct {
636+
Plan string `json:"plan"`
637+
Windows []usageWindow `json:"windows"`
638+
}
639+
640+
windows := make([]usageWindow, len(usage.Windows))
641+
for i, w := range usage.Windows {
642+
windows[i] = usageWindow{
643+
LimitID: w.LimitID,
644+
UsagePercentage: w.UsagePercentage,
645+
ResetsAt: time.Unix(w.ResetAt, 0).UTC().Format(time.RFC3339),
646+
WindowDurationMinutes: w.WindowDurationMinutes,
647+
}
648+
}
649+
650+
payload := usagePayload{
651+
Plan: usage.Plan,
652+
Windows: windows,
653+
}
654+
655+
err := model.SendHTTPRequestJSON(model.HTTPRequestOptions[usagePayload, any]{
656+
Context: ctx,
657+
Endpoint: model.Endpoint{
658+
Token: s.config.Token,
659+
APIEndpoint: s.config.APIEndpoint,
660+
},
661+
Method: "POST",
662+
Path: "/api/v1/codex-usage",
663+
Payload: payload,
664+
Timeout: 5 * time.Second,
665+
})
666+
if err != nil {
667+
slog.Warn("Failed to send codex usage to server", slog.Any("err", err))
668+
}
669+
}
670+
671+
// GetCachedCodexRateLimit returns a copy of the cached Codex rate limit data, or nil if not available.
672+
func (s *CCInfoTimerService) GetCachedCodexRateLimit() *CodexRateLimitData {
673+
s.codexRateLimitCache.mu.RLock()
674+
defer s.codexRateLimitCache.mu.RUnlock()
675+
676+
if s.codexRateLimitCache.usage == nil {
677+
return nil
678+
}
679+
680+
// Return a copy
681+
copy := *s.codexRateLimitCache.usage
682+
windowsCopy := make([]CodexRateLimitWindow, len(copy.Windows))
683+
for i, w := range copy.Windows {
684+
windowsCopy[i] = w
685+
}
686+
copy.Windows = windowsCopy
687+
return &copy
688+
}
689+
690+
// GetCachedCodexRateLimitError returns the last error from Codex rate limit fetching, or empty string if none.
691+
func (s *CCInfoTimerService) GetCachedCodexRateLimitError() string {
692+
s.codexRateLimitCache.mu.RLock()
693+
defer s.codexRateLimitCache.mu.RUnlock()
694+
return s.codexRateLimitCache.lastError
695+
}
696+
554697
// shortenAPIError converts an Anthropic usage API error into a short string for statusline display.
555698
func shortenAPIError(err error) string {
556699
msg := err.Error()

daemon/codex_ratelimit.go

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
package daemon
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"fmt"
7+
"net/http"
8+
"os"
9+
"path/filepath"
10+
"sync"
11+
"time"
12+
)
13+
14+
const codexUsageCacheTTL = 10 * time.Minute
15+
16+
// CodexRateLimitData holds the parsed rate limit data from the Codex API
17+
type CodexRateLimitData struct {
18+
Plan string
19+
Windows []CodexRateLimitWindow
20+
}
21+
22+
// CodexRateLimitWindow holds a single rate limit window from the Codex API
23+
type CodexRateLimitWindow struct {
24+
LimitID string
25+
UsagePercentage float64
26+
ResetAt int64 // Unix timestamp
27+
WindowDurationMinutes int
28+
}
29+
30+
type codexRateLimitCache struct {
31+
mu sync.RWMutex
32+
usage *CodexRateLimitData
33+
fetchedAt time.Time
34+
lastAttemptAt time.Time
35+
lastError string // short error description for statusline display
36+
}
37+
38+
// codexAuthData maps the relevant fields from ~/.codex/auth.json
39+
type codexAuthData struct {
40+
AccessToken string
41+
AccountID string
42+
}
43+
44+
// codexAuthJSON maps the full ~/.codex/auth.json structure
45+
type codexAuthJSON struct {
46+
AuthMode string `json:"authMode"`
47+
APIKey *string `json:"apiKey"`
48+
TokenData *codexAuthTokenData `json:"tokenData"`
49+
}
50+
51+
type codexAuthTokenData struct {
52+
AccessToken string `json:"accessToken"`
53+
RefreshToken string `json:"refreshToken"`
54+
IDTokenClaims *codexIDTokenClaims `json:"idTokenClaims"`
55+
}
56+
57+
type codexIDTokenClaims struct {
58+
AccountID string `json:"accountId"`
59+
}
60+
61+
// loadCodexAuth reads the Codex authentication data from ~/.codex/auth.json.
62+
func loadCodexAuth() (*codexAuthData, error) {
63+
homeDir, err := os.UserHomeDir()
64+
if err != nil {
65+
return nil, fmt.Errorf("failed to get home directory: %w", err)
66+
}
67+
68+
data, err := os.ReadFile(filepath.Join(homeDir, ".codex", "auth.json"))
69+
if err != nil {
70+
return nil, fmt.Errorf("codex auth file read failed: %w", err)
71+
}
72+
73+
var auth codexAuthJSON
74+
if err := json.Unmarshal(data, &auth); err != nil {
75+
return nil, fmt.Errorf("failed to parse codex auth JSON: %w", err)
76+
}
77+
78+
if auth.TokenData == nil || auth.TokenData.AccessToken == "" {
79+
return nil, fmt.Errorf("no access token found in codex auth")
80+
}
81+
82+
accountID := ""
83+
if auth.TokenData.IDTokenClaims != nil {
84+
accountID = auth.TokenData.IDTokenClaims.AccountID
85+
}
86+
87+
return &codexAuthData{
88+
AccessToken: auth.TokenData.AccessToken,
89+
AccountID: accountID,
90+
}, nil
91+
}
92+
93+
// codexUsageResponse maps the Codex usage API response
94+
type codexUsageResponse struct {
95+
RateLimits codexRateLimitSnapshot `json:"rateLimits"`
96+
}
97+
98+
type codexRateLimitSnapshot struct {
99+
Plan string `json:"plan"`
100+
RateLimitWindows []codexRateLimitWindowRaw `json:"rateLimitWindows"`
101+
}
102+
103+
type codexRateLimitWindowRaw struct {
104+
LimitID string `json:"limitId"`
105+
UsagePercentage float64 `json:"usagePercentage"`
106+
ResetAt int64 `json:"resetAt"`
107+
WindowDurationMinutes int `json:"windowDurationMinutes"`
108+
}
109+
110+
// fetchCodexUsage calls the Codex usage API and returns rate limit data.
111+
func fetchCodexUsage(ctx context.Context, auth *codexAuthData) (*CodexRateLimitData, error) {
112+
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://api.openai.com/api/codex/usage", nil)
113+
if err != nil {
114+
return nil, err
115+
}
116+
117+
req.Header.Set("Authorization", "Bearer "+auth.AccessToken)
118+
if auth.AccountID != "" {
119+
req.Header.Set("ChatGPT-Account-Id", auth.AccountID)
120+
}
121+
req.Header.Set("User-Agent", "shelltime-daemon")
122+
123+
client := &http.Client{Timeout: 5 * time.Second}
124+
resp, err := client.Do(req)
125+
if err != nil {
126+
return nil, err
127+
}
128+
defer resp.Body.Close()
129+
130+
if resp.StatusCode != http.StatusOK {
131+
return nil, fmt.Errorf("codex usage API returned status %d", resp.StatusCode)
132+
}
133+
134+
var usage codexUsageResponse
135+
if err := json.NewDecoder(resp.Body).Decode(&usage); err != nil {
136+
return nil, fmt.Errorf("failed to decode codex usage response: %w", err)
137+
}
138+
139+
windows := make([]CodexRateLimitWindow, len(usage.RateLimits.RateLimitWindows))
140+
for i, w := range usage.RateLimits.RateLimitWindows {
141+
windows[i] = CodexRateLimitWindow{
142+
LimitID: w.LimitID,
143+
UsagePercentage: w.UsagePercentage,
144+
ResetAt: w.ResetAt,
145+
WindowDurationMinutes: w.WindowDurationMinutes,
146+
}
147+
}
148+
149+
return &CodexRateLimitData{
150+
Plan: usage.RateLimits.Plan,
151+
Windows: windows,
152+
}, nil
153+
}
154+
155+
// shortenCodexAPIError converts a Codex usage API error into a short string for statusline display.
156+
func shortenCodexAPIError(err error) string {
157+
msg := err.Error()
158+
159+
var status int
160+
if _, scanErr := fmt.Sscanf(msg, "codex usage API returned status %d", &status); scanErr == nil {
161+
return fmt.Sprintf("api:%d", status)
162+
}
163+
164+
if len(msg) >= 6 && msg[:6] == "failed" {
165+
return "api:decode"
166+
}
167+
168+
return "network"
169+
}

0 commit comments

Comments
 (0)