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
19 changes: 19 additions & 0 deletions cmd/daemon/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,25 @@ func main() {
}
}

codexInstalled, err := daemon.CodexInstallationStatus()
if err != nil {
if reason, ok := daemon.CodexSyncSkipReason(err); ok {
slog.Info("Skipping Codex usage sync service startup", slog.String("reason", reason))
} else {
slog.Error("Failed to check Codex installation status", slog.Any("err", err))
}
} else if !codexInstalled {
slog.Info("Skipping Codex usage sync service startup", slog.String("reason", "codex_not_configured"))
} else {
codexUsageSyncService := daemon.NewCodexUsageSyncService(cfg)
if err := codexUsageSyncService.Start(ctx); err != nil {
slog.Error("Failed to start Codex usage sync service", slog.Any("err", err))
} else {
slog.Info("Codex usage sync service started")
defer codexUsageSyncService.Stop()
}
}

// Create processor instance
processor := daemon.NewSocketHandler(&cfg, pubsub)

Expand Down
145 changes: 1 addition & 144 deletions daemon/cc_info_timer.go
Original file line number Diff line number Diff line change
Expand Up @@ -55,9 +55,6 @@ 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
Expand All @@ -70,8 +67,7 @@ func NewCCInfoTimerService(config *model.ShellTimeConfig) *CCInfoTimerService {
cache: make(map[CCInfoTimeRange]CCInfoCache),
activeRanges: make(map[CCInfoTimeRange]bool),
gitCache: make(map[string]*GitCacheEntry),
rateLimitCache: &anthropicRateLimitCache{},
codexRateLimitCache: &codexRateLimitCache{},
rateLimitCache: &anthropicRateLimitCache{},
stopChan: make(chan struct{}),
}
}
Expand Down Expand Up @@ -156,11 +152,6 @@ 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")
}
Expand All @@ -180,7 +171,6 @@ 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())

Expand All @@ -204,7 +194,6 @@ func (s *CCInfoTimerService) timerLoop() {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
s.fetchRateLimit(ctx)
s.fetchCodexRateLimit(ctx)
}()

case <-s.stopChan:
Expand Down Expand Up @@ -562,138 +551,6 @@ 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 &copy
}

// 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()
Expand Down
104 changes: 100 additions & 4 deletions daemon/codex_ratelimit.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package daemon
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"os"
Expand All @@ -13,6 +14,19 @@ import (

const codexUsageCacheTTL = 10 * time.Minute

var (
loadCodexAuthFunc = loadCodexAuth
fetchCodexUsageFunc = fetchCodexUsage
codexPathExistsFunc = codexPathExists
)

var (
errCodexDirMissing = errors.New("codex directory missing")
errCodexAuthFileMissing = errors.New("codex auth file missing")
errCodexAuthInvalid = errors.New("codex auth invalid")
errCodexTokenInvalid = errors.New("codex token invalid")
)

// CodexRateLimitData holds the parsed rate limit data from the Codex API
type CodexRateLimitData struct {
Plan string
Expand Down Expand Up @@ -58,15 +72,79 @@ type codexIDTokenClaims struct {
AccountID string `json:"accountId"`
}

func codexConfigDirPath() (string, error) {
homeDir, err := os.UserHomeDir()
if err != nil {
return "", fmt.Errorf("failed to get home directory: %w", err)
}

return filepath.Join(homeDir, ".codex"), nil
}

func codexAuthFilePath() (string, error) {
dir, err := codexConfigDirPath()
if err != nil {
return "", err
}

return filepath.Join(dir, "auth.json"), nil
}

func codexPathExists(path string) (bool, error) {
_, err := os.Stat(path)
if err == nil {
return true, nil
}
if errors.Is(err, os.ErrNotExist) {
return false, nil
}
return false, err
}

func codexInstallationStatus() (bool, error) {
dirPath, err := codexConfigDirPath()
if err != nil {
return false, err
}
exists, err := codexPathExistsFunc(dirPath)
if err != nil {
return false, err
}
if !exists {
return false, errCodexDirMissing
}

authPath, err := codexAuthFilePath()
if err != nil {
return false, err
}
exists, err = codexPathExistsFunc(authPath)
if err != nil {
return false, err
}
if !exists {
return false, errCodexAuthFileMissing
}

return true, nil
}

func CodexInstallationStatus() (bool, error) {
return codexInstallationStatus()
}

// loadCodexAuth reads the Codex authentication data from ~/.codex/auth.json.
func loadCodexAuth() (*codexAuthData, error) {
homeDir, err := os.UserHomeDir()
authPath, err := codexAuthFilePath()
if err != nil {
return nil, fmt.Errorf("failed to get home directory: %w", err)
return nil, err
}

data, err := os.ReadFile(filepath.Join(homeDir, ".codex", "auth.json"))
data, err := os.ReadFile(authPath)
if err != nil {
if errors.Is(err, os.ErrNotExist) {
return nil, errCodexAuthFileMissing
}
return nil, fmt.Errorf("codex auth file read failed: %w", err)
}

Expand All @@ -76,7 +154,7 @@ func loadCodexAuth() (*codexAuthData, error) {
}

if auth.TokenData == nil || auth.TokenData.AccessToken == "" {
return nil, fmt.Errorf("no access token found in codex auth")
return nil, errCodexAuthInvalid
}

accountID := ""
Expand Down Expand Up @@ -128,6 +206,9 @@ func fetchCodexUsage(ctx context.Context, auth *codexAuthData) (*CodexRateLimitD
defer resp.Body.Close()

if resp.StatusCode != http.StatusOK {
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
return nil, errCodexTokenInvalid
}
return nil, fmt.Errorf("codex usage API returned status %d", resp.StatusCode)
}

Expand Down Expand Up @@ -167,3 +248,18 @@ func shortenCodexAPIError(err error) string {

return "network"
}

func CodexSyncSkipReason(err error) (string, bool) {
switch {
case errors.Is(err, errCodexDirMissing):
return "missing_codex_dir", true
case errors.Is(err, errCodexAuthFileMissing):
return "missing_auth_file", true
case errors.Is(err, errCodexAuthInvalid):
return "invalid_auth", true
case errors.Is(err, errCodexTokenInvalid):
return "invalid_auth_token", true
default:
return "", false
}
}
Loading
Loading