diff --git a/daemon/cc_info_timer.go b/daemon/cc_info_timer.go index 448defd..94d40fa 100644 --- a/daemon/cc_info_timer.go +++ b/daemon/cc_info_timer.go @@ -55,6 +55,9 @@ type CCInfoTimerService struct { // Anthropic rate limit cache rateLimitCache *anthropicRateLimitCache + // Codex rate limit cache + codexRateLimitCache *codexRateLimitCache + // User profile cache (permanent for daemon lifetime) userLogin string userLoginFetched bool @@ -67,7 +70,8 @@ func NewCCInfoTimerService(config *model.ShellTimeConfig) *CCInfoTimerService { cache: make(map[CCInfoTimeRange]CCInfoCache), activeRanges: make(map[CCInfoTimeRange]bool), gitCache: make(map[string]*GitCacheEntry), - rateLimitCache: &anthropicRateLimitCache{}, + rateLimitCache: &anthropicRateLimitCache{}, + codexRateLimitCache: &codexRateLimitCache{}, stopChan: make(chan struct{}), } } @@ -152,6 +156,11 @@ func (s *CCInfoTimerService) stopTimer() { s.rateLimitCache.fetchedAt = time.Time{} s.rateLimitCache.lastAttemptAt = time.Time{} s.rateLimitCache.mu.Unlock() + s.codexRateLimitCache.mu.Lock() + s.codexRateLimitCache.usage = nil + s.codexRateLimitCache.fetchedAt = time.Time{} + s.codexRateLimitCache.lastAttemptAt = time.Time{} + s.codexRateLimitCache.mu.Unlock() slog.Info("CC info timer stopped due to inactivity") } @@ -171,6 +180,7 @@ func (s *CCInfoTimerService) timerLoop() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() s.fetchRateLimit(ctx) + s.fetchCodexRateLimit(ctx) }() go s.fetchUserProfile(context.Background()) @@ -194,6 +204,7 @@ func (s *CCInfoTimerService) timerLoop() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() s.fetchRateLimit(ctx) + s.fetchCodexRateLimit(ctx) }() case <-s.stopChan: @@ -551,6 +562,138 @@ func (s *CCInfoTimerService) GetCachedRateLimitError() string { return s.rateLimitCache.lastError } +// fetchCodexRateLimit fetches Codex rate limit data if cache is stale. +func (s *CCInfoTimerService) fetchCodexRateLimit(ctx context.Context) { + if runtime.GOOS != "darwin" && runtime.GOOS != "linux" { + return + } + + // Check cache TTL under read lock + s.codexRateLimitCache.mu.RLock() + sinceLastFetch := time.Since(s.codexRateLimitCache.fetchedAt) + sinceLastAttempt := time.Since(s.codexRateLimitCache.lastAttemptAt) + s.codexRateLimitCache.mu.RUnlock() + + if sinceLastFetch < codexUsageCacheTTL || sinceLastAttempt < codexUsageCacheTTL { + return + } + + // Record attempt time + s.codexRateLimitCache.mu.Lock() + s.codexRateLimitCache.lastAttemptAt = time.Now() + s.codexRateLimitCache.mu.Unlock() + + auth, err := loadCodexAuth() + if err != nil || auth == nil { + slog.Debug("Failed to load Codex auth", slog.Any("err", err)) + s.codexRateLimitCache.mu.Lock() + s.codexRateLimitCache.lastError = "auth" + s.codexRateLimitCache.mu.Unlock() + return + } + + usage, err := fetchCodexUsage(ctx, auth) + if err != nil { + slog.Warn("Failed to fetch Codex usage", slog.Any("err", err)) + s.codexRateLimitCache.mu.Lock() + s.codexRateLimitCache.lastError = shortenCodexAPIError(err) + s.codexRateLimitCache.mu.Unlock() + return + } + + s.codexRateLimitCache.mu.Lock() + s.codexRateLimitCache.usage = usage + s.codexRateLimitCache.fetchedAt = time.Now() + s.codexRateLimitCache.lastError = "" + s.codexRateLimitCache.mu.Unlock() + + // Send usage data to server (fire-and-forget) + go func() { + bgCtx, bgCancel := context.WithTimeout(context.Background(), 10*time.Second) + defer bgCancel() + s.sendCodexUsageToServer(bgCtx, usage) + }() + + slog.Debug("Codex rate limit updated", + slog.String("plan", usage.Plan), + slog.Int("windows", len(usage.Windows))) +} + +// sendCodexUsageToServer sends Codex usage data to the ShellTime server +// for scheduling push notifications when rate limits reset. +func (s *CCInfoTimerService) sendCodexUsageToServer(ctx context.Context, usage *CodexRateLimitData) { + if s.config.Token == "" { + return + } + + type usageWindow struct { + LimitID string `json:"limit_id"` + UsagePercentage float64 `json:"usage_percentage"` + ResetsAt string `json:"resets_at"` + WindowDurationMinutes int `json:"window_duration_minutes"` + } + type usagePayload struct { + Plan string `json:"plan"` + Windows []usageWindow `json:"windows"` + } + + windows := make([]usageWindow, len(usage.Windows)) + for i, w := range usage.Windows { + windows[i] = usageWindow{ + LimitID: w.LimitID, + UsagePercentage: w.UsagePercentage, + ResetsAt: time.Unix(w.ResetAt, 0).UTC().Format(time.RFC3339), + WindowDurationMinutes: w.WindowDurationMinutes, + } + } + + payload := usagePayload{ + Plan: usage.Plan, + Windows: windows, + } + + err := model.SendHTTPRequestJSON(model.HTTPRequestOptions[usagePayload, any]{ + Context: ctx, + Endpoint: model.Endpoint{ + Token: s.config.Token, + APIEndpoint: s.config.APIEndpoint, + }, + Method: "POST", + Path: "/api/v1/codex-usage", + Payload: payload, + Timeout: 5 * time.Second, + }) + if err != nil { + slog.Warn("Failed to send codex usage to server", slog.Any("err", err)) + } +} + +// GetCachedCodexRateLimit returns a copy of the cached Codex rate limit data, or nil if not available. +func (s *CCInfoTimerService) GetCachedCodexRateLimit() *CodexRateLimitData { + s.codexRateLimitCache.mu.RLock() + defer s.codexRateLimitCache.mu.RUnlock() + + if s.codexRateLimitCache.usage == nil { + return nil + } + + // Return a copy + copy := *s.codexRateLimitCache.usage + windowsCopy := make([]CodexRateLimitWindow, len(copy.Windows)) + for i, w := range copy.Windows { + windowsCopy[i] = w + } + copy.Windows = windowsCopy + return © +} + +// GetCachedCodexRateLimitError returns the last error from Codex rate limit fetching, or empty string if none. +func (s *CCInfoTimerService) GetCachedCodexRateLimitError() string { + s.codexRateLimitCache.mu.RLock() + defer s.codexRateLimitCache.mu.RUnlock() + return s.codexRateLimitCache.lastError +} + // shortenAPIError converts an Anthropic usage API error into a short string for statusline display. func shortenAPIError(err error) string { msg := err.Error() diff --git a/daemon/codex_ratelimit.go b/daemon/codex_ratelimit.go new file mode 100644 index 0000000..127171f --- /dev/null +++ b/daemon/codex_ratelimit.go @@ -0,0 +1,169 @@ +package daemon + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "os" + "path/filepath" + "sync" + "time" +) + +const codexUsageCacheTTL = 10 * time.Minute + +// CodexRateLimitData holds the parsed rate limit data from the Codex API +type CodexRateLimitData struct { + Plan string + Windows []CodexRateLimitWindow +} + +// CodexRateLimitWindow holds a single rate limit window from the Codex API +type CodexRateLimitWindow struct { + LimitID string + UsagePercentage float64 + ResetAt int64 // Unix timestamp + WindowDurationMinutes int +} + +type codexRateLimitCache struct { + mu sync.RWMutex + usage *CodexRateLimitData + fetchedAt time.Time + lastAttemptAt time.Time + lastError string // short error description for statusline display +} + +// codexAuthData maps the relevant fields from ~/.codex/auth.json +type codexAuthData struct { + AccessToken string + AccountID string +} + +// codexAuthJSON maps the full ~/.codex/auth.json structure +type codexAuthJSON struct { + AuthMode string `json:"authMode"` + APIKey *string `json:"apiKey"` + TokenData *codexAuthTokenData `json:"tokenData"` +} + +type codexAuthTokenData struct { + AccessToken string `json:"accessToken"` + RefreshToken string `json:"refreshToken"` + IDTokenClaims *codexIDTokenClaims `json:"idTokenClaims"` +} + +type codexIDTokenClaims struct { + AccountID string `json:"accountId"` +} + +// loadCodexAuth reads the Codex authentication data from ~/.codex/auth.json. +func loadCodexAuth() (*codexAuthData, error) { + homeDir, err := os.UserHomeDir() + if err != nil { + return nil, fmt.Errorf("failed to get home directory: %w", err) + } + + data, err := os.ReadFile(filepath.Join(homeDir, ".codex", "auth.json")) + if err != nil { + return nil, fmt.Errorf("codex auth file read failed: %w", err) + } + + var auth codexAuthJSON + if err := json.Unmarshal(data, &auth); err != nil { + return nil, fmt.Errorf("failed to parse codex auth JSON: %w", err) + } + + if auth.TokenData == nil || auth.TokenData.AccessToken == "" { + return nil, fmt.Errorf("no access token found in codex auth") + } + + accountID := "" + if auth.TokenData.IDTokenClaims != nil { + accountID = auth.TokenData.IDTokenClaims.AccountID + } + + return &codexAuthData{ + AccessToken: auth.TokenData.AccessToken, + AccountID: accountID, + }, nil +} + +// codexUsageResponse maps the Codex usage API response +type codexUsageResponse struct { + RateLimits codexRateLimitSnapshot `json:"rateLimits"` +} + +type codexRateLimitSnapshot struct { + Plan string `json:"plan"` + RateLimitWindows []codexRateLimitWindowRaw `json:"rateLimitWindows"` +} + +type codexRateLimitWindowRaw struct { + LimitID string `json:"limitId"` + UsagePercentage float64 `json:"usagePercentage"` + ResetAt int64 `json:"resetAt"` + WindowDurationMinutes int `json:"windowDurationMinutes"` +} + +// fetchCodexUsage calls the Codex usage API and returns rate limit data. +func fetchCodexUsage(ctx context.Context, auth *codexAuthData) (*CodexRateLimitData, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://api.openai.com/api/codex/usage", nil) + if err != nil { + return nil, err + } + + req.Header.Set("Authorization", "Bearer "+auth.AccessToken) + if auth.AccountID != "" { + req.Header.Set("ChatGPT-Account-Id", auth.AccountID) + } + req.Header.Set("User-Agent", "shelltime-daemon") + + client := &http.Client{Timeout: 5 * time.Second} + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("codex usage API returned status %d", resp.StatusCode) + } + + var usage codexUsageResponse + if err := json.NewDecoder(resp.Body).Decode(&usage); err != nil { + return nil, fmt.Errorf("failed to decode codex usage response: %w", err) + } + + windows := make([]CodexRateLimitWindow, len(usage.RateLimits.RateLimitWindows)) + for i, w := range usage.RateLimits.RateLimitWindows { + windows[i] = CodexRateLimitWindow{ + LimitID: w.LimitID, + UsagePercentage: w.UsagePercentage, + ResetAt: w.ResetAt, + WindowDurationMinutes: w.WindowDurationMinutes, + } + } + + return &CodexRateLimitData{ + Plan: usage.RateLimits.Plan, + Windows: windows, + }, nil +} + +// shortenCodexAPIError converts a Codex usage API error into a short string for statusline display. +func shortenCodexAPIError(err error) string { + msg := err.Error() + + var status int + if _, scanErr := fmt.Sscanf(msg, "codex usage API returned status %d", &status); scanErr == nil { + return fmt.Sprintf("api:%d", status) + } + + if len(msg) >= 6 && msg[:6] == "failed" { + return "api:decode" + } + + return "network" +}