From 77cd44be2e2e9a66ccbaba4f5fbc95207056099f Mon Sep 17 00:00:00 2001 From: AnnatarHe Date: Mon, 6 Apr 2026 23:27:16 +0800 Subject: [PATCH 1/2] fix(daemon): decouple codex usage sync from cc statusline --- cmd/daemon/main.go | 8 ++ daemon/cc_info_timer.go | 145 +----------------------------- daemon/codex_ratelimit.go | 5 ++ daemon/codex_usage_sync.go | 145 ++++++++++++++++++++++++++++++ daemon/codex_usage_sync_test.go | 153 ++++++++++++++++++++++++++++++++ 5 files changed, 312 insertions(+), 144 deletions(-) create mode 100644 daemon/codex_usage_sync.go create mode 100644 daemon/codex_usage_sync_test.go diff --git a/cmd/daemon/main.go b/cmd/daemon/main.go index abfd96a..fa322b9 100644 --- a/cmd/daemon/main.go +++ b/cmd/daemon/main.go @@ -148,6 +148,14 @@ func main() { } } + 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) diff --git a/daemon/cc_info_timer.go b/daemon/cc_info_timer.go index 94d40fa..448defd 100644 --- a/daemon/cc_info_timer.go +++ b/daemon/cc_info_timer.go @@ -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 @@ -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{}), } } @@ -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") } @@ -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()) @@ -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: @@ -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 © -} - -// 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 index 127171f..dc39cb3 100644 --- a/daemon/codex_ratelimit.go +++ b/daemon/codex_ratelimit.go @@ -13,6 +13,11 @@ import ( const codexUsageCacheTTL = 10 * time.Minute +var ( + loadCodexAuthFunc = loadCodexAuth + fetchCodexUsageFunc = fetchCodexUsage +) + // CodexRateLimitData holds the parsed rate limit data from the Codex API type CodexRateLimitData struct { Plan string diff --git a/daemon/codex_usage_sync.go b/daemon/codex_usage_sync.go new file mode 100644 index 0000000..0990daa --- /dev/null +++ b/daemon/codex_usage_sync.go @@ -0,0 +1,145 @@ +package daemon + +import ( + "context" + "fmt" + "log/slog" + "sync" + "time" + + "github.com/malamtime/cli/model" +) + +var CodexUsageSyncInterval = 10 * time.Minute + +// CodexUsageSyncService periodically fetches Codex usage and syncs it to the server. +type CodexUsageSyncService struct { + config model.ShellTimeConfig + ticker *time.Ticker + stopChan chan struct{} + wg sync.WaitGroup +} + +// NewCodexUsageSyncService creates a new Codex usage sync service. +func NewCodexUsageSyncService(config model.ShellTimeConfig) *CodexUsageSyncService { + return &CodexUsageSyncService{ + config: config, + stopChan: make(chan struct{}), + } +} + +// Start begins the periodic Codex usage sync job. +func (s *CodexUsageSyncService) Start(ctx context.Context) error { + s.ticker = time.NewTicker(CodexUsageSyncInterval) + s.wg.Add(1) + + go func() { + defer s.wg.Done() + + s.sync() + + for { + select { + case <-s.ticker.C: + s.sync() + case <-s.stopChan: + return + case <-ctx.Done(): + return + } + } + }() + + slog.Info("Codex usage sync service started", slog.Duration("interval", CodexUsageSyncInterval)) + return nil +} + +// Stop stops the Codex usage sync service. +func (s *CodexUsageSyncService) Stop() { + if s.ticker != nil { + s.ticker.Stop() + } + close(s.stopChan) + s.wg.Wait() + slog.Info("Codex usage sync service stopped") +} + +func (s *CodexUsageSyncService) sync() { + if s.config.Token == "" { + return + } + + if err := syncCodexUsage(context.Background(), s.config); err != nil { + slog.Warn("Failed to sync codex usage", slog.Any("err", err)) + } +} + +func syncCodexUsage(ctx context.Context, config model.ShellTimeConfig) error { + if config.Token == "" { + return nil + } + + runCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + + auth, err := loadCodexAuthFunc() + if err != nil || auth == nil { + if err == nil && auth == nil { + err = fmt.Errorf("codex auth unavailable") + } + return err + } + + usage, err := fetchCodexUsageFunc(runCtx, auth) + if err != nil { + return err + } + + return sendCodexUsageToServer(runCtx, config, usage) +} + +// sendCodexUsageToServer sends Codex usage data to the ShellTime server +// for scheduling push notifications when rate limits reset. +func sendCodexUsageToServer(ctx context.Context, config model.ShellTimeConfig, usage *CodexRateLimitData) error { + if config.Token == "" { + return nil + } + + 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, + } + + return model.SendHTTPRequestJSON(model.HTTPRequestOptions[usagePayload, any]{ + Context: ctx, + Endpoint: model.Endpoint{ + Token: config.Token, + APIEndpoint: config.APIEndpoint, + }, + Method: "POST", + Path: "/api/v1/codex-usage", + Payload: payload, + Timeout: 5 * time.Second, + }) +} diff --git a/daemon/codex_usage_sync_test.go b/daemon/codex_usage_sync_test.go new file mode 100644 index 0000000..2b44b9c --- /dev/null +++ b/daemon/codex_usage_sync_test.go @@ -0,0 +1,153 @@ +package daemon + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/malamtime/cli/model" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSyncCodexUsage_SendsUsageToServer(t *testing.T) { + t.Helper() + + originalLoad := loadCodexAuthFunc + originalFetch := fetchCodexUsageFunc + defer func() { + loadCodexAuthFunc = originalLoad + fetchCodexUsageFunc = originalFetch + }() + + loadCodexAuthFunc = func() (*codexAuthData, error) { + return &codexAuthData{AccessToken: "test-token", AccountID: "acct-1"}, nil + } + fetchCodexUsageFunc = func(ctx context.Context, auth *codexAuthData) (*CodexRateLimitData, error) { + return &CodexRateLimitData{ + Plan: "pro", + Windows: []CodexRateLimitWindow{ + { + LimitID: "main", + UsagePercentage: 72.5, + ResetAt: 1712400000, + WindowDurationMinutes: 300, + }, + }, + }, nil + } + + var captured struct { + Plan string `json:"plan"` + Windows []struct { + LimitID string `json:"limit_id"` + UsagePercentage float64 `json:"usage_percentage"` + ResetsAt string `json:"resets_at"` + WindowDurationMinutes int `json:"window_duration_minutes"` + } `json:"windows"` + } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/api/v1/codex-usage", r.URL.Path) + assert.Equal(t, "CLI shelltime-token", r.Header.Get("Authorization")) + require.NoError(t, json.NewDecoder(r.Body).Decode(&captured)) + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + cfg := model.ShellTimeConfig{ + Token: "shelltime-token", + APIEndpoint: server.URL, + } + + err := syncCodexUsage(context.Background(), cfg) + require.NoError(t, err) + + require.Len(t, captured.Windows, 1) + assert.Equal(t, "pro", captured.Plan) + assert.Equal(t, "main", captured.Windows[0].LimitID) + assert.Equal(t, 72.5, captured.Windows[0].UsagePercentage) + assert.Equal(t, "2024-04-06T10:40:00Z", captured.Windows[0].ResetsAt) + assert.Equal(t, 300, captured.Windows[0].WindowDurationMinutes) +} + +func TestSyncCodexUsage_AuthError(t *testing.T) { + t.Helper() + + originalLoad := loadCodexAuthFunc + originalFetch := fetchCodexUsageFunc + defer func() { + loadCodexAuthFunc = originalLoad + fetchCodexUsageFunc = originalFetch + }() + + loadCodexAuthFunc = func() (*codexAuthData, error) { + return nil, assert.AnError + } + fetchCodexUsageFunc = func(ctx context.Context, auth *codexAuthData) (*CodexRateLimitData, error) { + t.Fatal("fetchCodexUsageFunc should not be called when auth loading fails") + return nil, nil + } + + cfg := model.ShellTimeConfig{Token: "shelltime-token"} + err := syncCodexUsage(context.Background(), cfg) + require.Error(t, err) + assert.ErrorIs(t, err, assert.AnError) +} + +func TestCodexUsageSyncService_StartRunsImmediatelyAndOnTicker(t *testing.T) { + t.Helper() + + originalInterval := CodexUsageSyncInterval + originalLoad := loadCodexAuthFunc + originalFetch := fetchCodexUsageFunc + defer func() { + CodexUsageSyncInterval = originalInterval + loadCodexAuthFunc = originalLoad + fetchCodexUsageFunc = originalFetch + }() + + CodexUsageSyncInterval = 20 * time.Millisecond + + loadCodexAuthFunc = func() (*codexAuthData, error) { + return &codexAuthData{AccessToken: "test-token"}, nil + } + fetchCodexUsageFunc = func(ctx context.Context, auth *codexAuthData) (*CodexRateLimitData, error) { + return &CodexRateLimitData{ + Plan: "pro", + Windows: []CodexRateLimitWindow{ + {LimitID: "main", UsagePercentage: 12, ResetAt: time.Now().Unix(), WindowDurationMinutes: 300}, + }, + }, nil + } + + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + service := NewCodexUsageSyncService(model.ShellTimeConfig{ + Token: "shelltime-token", + APIEndpoint: server.URL, + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + require.NoError(t, service.Start(ctx)) + + require.Eventually(t, func() bool { + return calls.Load() >= 2 + }, 250*time.Millisecond, 10*time.Millisecond) + + service.Stop() + stoppedAt := calls.Load() + time.Sleep(50 * time.Millisecond) + assert.Equal(t, stoppedAt, calls.Load()) +} From d1b334b2f1dd5ab863c374a23a343ad3a6fe90da Mon Sep 17 00:00:00 2001 From: AnnatarHe Date: Mon, 6 Apr 2026 23:40:13 +0800 Subject: [PATCH 2/2] fix(daemon): skip codex sync when auth is unavailable --- cmd/daemon/main.go | 21 +++++-- daemon/codex_ratelimit.go | 99 +++++++++++++++++++++++++++++-- daemon/codex_usage_sync.go | 13 +++-- daemon/codex_usage_sync_test.go | 100 ++++++++++++++++++++++++++++++++ 4 files changed, 219 insertions(+), 14 deletions(-) diff --git a/cmd/daemon/main.go b/cmd/daemon/main.go index fa322b9..a9280dd 100644 --- a/cmd/daemon/main.go +++ b/cmd/daemon/main.go @@ -148,12 +148,23 @@ func main() { } } - codexUsageSyncService := daemon.NewCodexUsageSyncService(cfg) - if err := codexUsageSyncService.Start(ctx); err != nil { - slog.Error("Failed to start Codex usage sync service", slog.Any("err", err)) + 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 { - slog.Info("Codex usage sync service started") - defer codexUsageSyncService.Stop() + 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 diff --git a/daemon/codex_ratelimit.go b/daemon/codex_ratelimit.go index dc39cb3..9fbfbe4 100644 --- a/daemon/codex_ratelimit.go +++ b/daemon/codex_ratelimit.go @@ -3,6 +3,7 @@ package daemon import ( "context" "encoding/json" + "errors" "fmt" "net/http" "os" @@ -16,6 +17,14 @@ 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 @@ -63,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) } @@ -81,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 := "" @@ -133,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) } @@ -172,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 + } +} diff --git a/daemon/codex_usage_sync.go b/daemon/codex_usage_sync.go index 0990daa..7214dc9 100644 --- a/daemon/codex_usage_sync.go +++ b/daemon/codex_usage_sync.go @@ -2,7 +2,6 @@ package daemon import ( "context" - "fmt" "log/slog" "sync" "time" @@ -70,6 +69,10 @@ func (s *CodexUsageSyncService) sync() { } if err := syncCodexUsage(context.Background(), s.config); err != nil { + if reason, ok := CodexSyncSkipReason(err); ok { + slog.Info("Skipping codex usage sync", slog.String("reason", reason)) + return + } slog.Warn("Failed to sync codex usage", slog.Any("err", err)) } } @@ -83,12 +86,12 @@ func syncCodexUsage(ctx context.Context, config model.ShellTimeConfig) error { defer cancel() auth, err := loadCodexAuthFunc() - if err != nil || auth == nil { - if err == nil && auth == nil { - err = fmt.Errorf("codex auth unavailable") - } + if err != nil { return err } + if auth == nil { + return errCodexAuthInvalid + } usage, err := fetchCodexUsageFunc(runCtx, auth) if err != nil { diff --git a/daemon/codex_usage_sync_test.go b/daemon/codex_usage_sync_test.go index 2b44b9c..4086a13 100644 --- a/daemon/codex_usage_sync_test.go +++ b/daemon/codex_usage_sync_test.go @@ -3,6 +3,7 @@ package daemon import ( "context" "encoding/json" + "errors" "net/http" "net/http/httptest" "sync/atomic" @@ -99,6 +100,105 @@ func TestSyncCodexUsage_AuthError(t *testing.T) { assert.ErrorIs(t, err, assert.AnError) } +func TestSyncCodexUsage_InvalidTokenSkips(t *testing.T) { + t.Helper() + + originalLoad := loadCodexAuthFunc + originalFetch := fetchCodexUsageFunc + defer func() { + loadCodexAuthFunc = originalLoad + fetchCodexUsageFunc = originalFetch + }() + + loadCodexAuthFunc = func() (*codexAuthData, error) { + return &codexAuthData{AccessToken: "test-token"}, nil + } + fetchCodexUsageFunc = func(ctx context.Context, auth *codexAuthData) (*CodexRateLimitData, error) { + return nil, errCodexTokenInvalid + } + + cfg := model.ShellTimeConfig{Token: "shelltime-token"} + err := syncCodexUsage(context.Background(), cfg) + require.Error(t, err) + + reason, ok := CodexSyncSkipReason(err) + require.True(t, ok) + assert.Equal(t, "invalid_auth_token", reason) +} + +func TestCodexInstallationStatus_MissingCodexDir(t *testing.T) { + originalExists := codexPathExistsFunc + defer func() { + codexPathExistsFunc = originalExists + }() + + codexPathExistsFunc = func(path string) (bool, error) { + return false, nil + } + + ok, err := CodexInstallationStatus() + assert.False(t, ok) + assert.ErrorIs(t, err, errCodexDirMissing) +} + +func TestCodexInstallationStatus_MissingAuthFile(t *testing.T) { + originalExists := codexPathExistsFunc + defer func() { + codexPathExistsFunc = originalExists + }() + + callCount := 0 + codexPathExistsFunc = func(path string) (bool, error) { + callCount++ + if callCount == 1 { + return true, nil + } + return false, nil + } + + ok, err := CodexInstallationStatus() + assert.False(t, ok) + assert.ErrorIs(t, err, errCodexAuthFileMissing) +} + +func TestCodexInstallationStatus_Ready(t *testing.T) { + originalExists := codexPathExistsFunc + defer func() { + codexPathExistsFunc = originalExists + }() + + codexPathExistsFunc = func(path string) (bool, error) { + return true, nil + } + + ok, err := CodexInstallationStatus() + require.NoError(t, err) + assert.True(t, ok) +} + +func TestCodexSyncSkipReason(t *testing.T) { + testCases := []struct { + name string + err error + expected string + ok bool + }{ + {name: "missing dir", err: errCodexDirMissing, expected: "missing_codex_dir", ok: true}, + {name: "missing auth file", err: errCodexAuthFileMissing, expected: "missing_auth_file", ok: true}, + {name: "invalid auth", err: errCodexAuthInvalid, expected: "invalid_auth", ok: true}, + {name: "invalid token", err: errCodexTokenInvalid, expected: "invalid_auth_token", ok: true}, + {name: "other error", err: errors.New("boom"), expected: "", ok: false}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + reason, ok := CodexSyncSkipReason(tc.err) + assert.Equal(t, tc.ok, ok) + assert.Equal(t, tc.expected, reason) + }) + } +} + func TestCodexUsageSyncService_StartRunsImmediatelyAndOnTicker(t *testing.T) { t.Helper()