diff --git a/README.md b/README.md index 52464b4..506cd5c 100644 --- a/README.md +++ b/README.md @@ -87,7 +87,7 @@ Server without making that account active: # Query the active managed account now. codex-switch account usage -# Query one saved account, or all accounts concurrently. +# Query one saved account, or all accounts in a reliable sequence. codex-switch account usage work codex-switch account usage --all @@ -100,10 +100,18 @@ codex-switch account usage work --cached ``` Normal `account list` calls refresh only missing snapshots or snapshots older -than 60 seconds. Each query runs in an isolated temporary `CODEX_HOME`; it does +than five minutes. Explicit `--refresh` queries every account. Failed queries +are retried up to three times with backoff, and multi-account refreshes run in a +reliable sequence instead of starting several Codex App Servers at once. Each +query runs in an isolated temporary `CODEX_HOME`; it does not switch `$CODEX_HOME/auth.json`, sessions, plugins, or UI state. The cache contains usage numbers and public account metadata only, never tokens. +The compact list preserves the last successful values when a refresh fails, +shows an error summary in the `ERROR` column, and prints the complete redacted +error below the table. The last failure is retained in the local usage cache so +`--cached` remains diagnostic after the failing process exits. + The compact list shows the next automatic reset for each quota window. Its `RESETS` column is separate: it reports earned manual rate-limit resets and the nearest known expiration, when Codex provides them. `account usage ` diff --git a/docs/architecture.md b/docs/architecture.md index 78f9998..9ae4dec 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -82,17 +82,20 @@ acquire shared operation lock -> create one temporary CODEX_HOME per profile -> write only that profile plus file-store config -> initialize the official Codex App Server - -> read account/rateLimits/read and account/usage/read + -> read account/read, account/rateLimits/read, and account/usage/read in order -> retain quota reset times and any earned-reset count and expiration metadata -> stop the server and delete the temporary home - -> reconcile any newer credential generation + -> reconcile any newer credential generation, including one produced by a failed query -> atomically save credential-free usage snapshots ``` The live account selection never changes. If Codex rotates a refresh token while answering the query, identity and generation checks run before the new document is saved. For an active profile, compare-before-replace protects the live -projection from a concurrent Codex write. +projection from a concurrent Codex write. Transient failures are retried with +backoff using the newest validated credential generation returned by the prior +attempt. Multi-account refreshes are intentionally sequential so independent +App Server processes do not compete for startup or authentication resources. ## Data locations diff --git a/docs/security.md b/docs/security.md index b59a940..56356a3 100644 --- a/docs/security.md +++ b/docs/security.md @@ -45,7 +45,9 @@ focuses on: never contains authentication documents. - A token refreshed during an isolated query is accepted only after account, workspace, and refresh-generation checks. Active-file updates use a - compare-before-replace check under the shared operation lock. + compare-before-replace check under the shared operation lock. A validated + newer generation is preserved even when the accompanying usage request fails, + so retries do not fall back to a rotated token. - Active-account detection uses the live account and workspace identifiers; stale recorded state, aliases, and email addresses cannot redirect live credentials into another saved profile. Synchronization rechecks the live diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index a785b48..68a0478 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -123,7 +123,11 @@ codex-switch account usage --cached ``` A stale cached snapshot is labeled `stale`; a failed refresh keeps that snapshot -and displays a warning rather than discarding useful data. +and displays the complete redacted error rather than discarding useful data. +The tool retries a failed isolated query up to three times with backoff and +queries saved accounts sequentially to avoid competing Codex App Server startup +or authentication work. The last failure remains visible with `--cached` until +that account refreshes successfully. ## Interrupted switch journal diff --git a/internal/accountusage/service.go b/internal/accountusage/service.go index be42fd5..7398473 100644 --- a/internal/accountusage/service.go +++ b/internal/accountusage/service.go @@ -7,7 +7,6 @@ import ( "fmt" "os" "path/filepath" - "sync" "time" "github.com/SilkageNet/codex-switch/internal/atomicfile" @@ -22,7 +21,11 @@ import ( "github.com/SilkageNet/codex-switch/internal/vault" ) -const queryTimeout = 20 * time.Second +const ( + queryTimeout = 20 * time.Second + defaultQueryAttempts = 3 + defaultRetryBaseDelay = 250 * time.Millisecond +) type QueryRunner interface { Query(context.Context, json.RawMessage) (codexusage.Snapshot, json.RawMessage, error) @@ -33,12 +36,16 @@ type Service struct { Paths appconfig.Paths Vault *vault.Manager Runner QueryRunner + + queryAttempts int + retryBaseDelay time.Duration } type Result struct { ProfileID string `json:"profileId"` Snapshot codexusage.Snapshot `json:"snapshot,omitempty"` Error string `json:"error,omitempty"` + Attempts int `json:"attempts,omitempty"` } type queryResult struct { @@ -46,6 +53,7 @@ type queryResult struct { snapshot codexusage.Snapshot auth json.RawMessage err error + attempts int } func (service Service) Cached() (usagecache.Cache, error) { @@ -79,29 +87,11 @@ func (service Service) Refresh(ctx context.Context, profileIDs []string) (map[st } } - queried := make(chan queryResult, len(selected)) - semaphore := make(chan struct{}, 4) - var group sync.WaitGroup + queried := make([]queryResult, 0, len(selected)) for _, profile := range selected { - profile := profile - group.Add(1) - go func() { - defer group.Done() - select { - case semaphore <- struct{}{}: - defer func() { <-semaphore }() - case <-ctx.Done(): - queried <- queryResult{profileID: profile.ID, err: ctx.Err()} - return - } - queryContext, cancel := context.WithTimeout(ctx, queryTimeout) - defer cancel() - snapshot, updatedAuth, queryErr := service.Runner.Query(queryContext, profile.Auth) - queried <- queryResult{profileID: profile.ID, snapshot: snapshot, auth: updatedAuth, err: queryErr} - }() + snapshot, updatedAuth, attempts, queryErr := service.queryWithRetry(ctx, profile.Auth) + queried = append(queried, queryResult{profileID: profile.ID, snapshot: snapshot, auth: updatedAuth, err: queryErr, attempts: attempts}) } - group.Wait() - close(queried) results := make(map[string]Result, len(selected)) candidates := make(map[string]json.RawMessage, len(selected)) @@ -109,12 +99,14 @@ func (service Service) Refresh(ctx context.Context, profileIDs []string) (map[st if err != nil { return nil, err } - for result := range queried { - entry := Result{ProfileID: result.profileID} + for _, result := range queried { + entry := Result{ProfileID: result.profileID, Attempts: result.attempts} if result.err != nil { entry.Error = result.err.Error() } else { entry.Snapshot = result.snapshot + } + if len(result.auth) > 0 { candidates[result.profileID] = result.auth } results[result.profileID] = entry @@ -159,7 +151,9 @@ func (service Service) Refresh(ctx context.Context, profileIDs []string) (map[st results[profileID] = entry continue } - cache.Profiles[profileID] = results[profileID].Snapshot + if results[profileID].Error == "" { + cache.Profiles[profileID] = results[profileID].Snapshot + } saved, parseErr := authschema.Parse(profile.Auth) if parseErr != nil { entry := results[profileID] @@ -210,12 +204,73 @@ func (service Service) Refresh(ctx context.Context, profileIDs []string) (map[st delete(cache.Profiles, profileID) } } + for profileID, result := range results { + if result.Error == "" { + delete(cache.Failures, profileID) + continue + } + cache.Failures[profileID] = usagecache.Failure{ + AttemptedAt: time.Now().UTC(), + Message: result.Error, + Attempts: result.Attempts, + } + } + for profileID := range cache.Failures { + if !validProfiles[profileID] { + delete(cache.Failures, profileID) + } + } if err := usagecache.Save(service.Paths.UsageCache, cache); err != nil { return nil, err } return results, nil } +func (service Service) queryWithRetry(ctx context.Context, auth json.RawMessage) (codexusage.Snapshot, json.RawMessage, int, error) { + attemptLimit := service.queryAttempts + if attemptLimit <= 0 { + attemptLimit = defaultQueryAttempts + } + baseDelay := service.retryBaseDelay + if baseDelay <= 0 { + baseDelay = defaultRetryBaseDelay + } + currentAuth := append(json.RawMessage(nil), auth...) + var latestAuth json.RawMessage + var lastErr error + for attempt := 1; attempt <= attemptLimit; attempt++ { + if err := ctx.Err(); err != nil { + return codexusage.Snapshot{}, nil, attempt - 1, err + } + queryContext, cancel := context.WithTimeout(ctx, queryTimeout) + snapshot, updatedAuth, err := service.Runner.Query(queryContext, currentAuth) + cancel() + if len(updatedAuth) > 0 { + latestAuth = append(latestAuth[:0], updatedAuth...) + currentAuth = append(currentAuth[:0], updatedAuth...) + } + if err == nil { + if len(updatedAuth) == 0 { + updatedAuth = latestAuth + } + return snapshot, updatedAuth, attempt, nil + } + lastErr = err + if attempt == attemptLimit { + break + } + delay := baseDelay << (attempt - 1) + timer := time.NewTimer(delay) + select { + case <-ctx.Done(): + timer.Stop() + return codexusage.Snapshot{}, latestAuth, attempt, ctx.Err() + case <-timer.C: + } + } + return codexusage.Snapshot{}, latestAuth, attemptLimit, fmt.Errorf("query failed after %d attempts: %w", attemptLimit, lastErr) +} + func (service Service) reconcileActive(profile *vault.Profile, candidate authschema.Document, state *appstate.State) (bool, error) { liveRaw, err := service.Home.ReadAuth() if errors.Is(err, os.ErrNotExist) { diff --git a/internal/accountusage/service_test.go b/internal/accountusage/service_test.go index 0dc7dbe..d1a331b 100644 --- a/internal/accountusage/service_test.go +++ b/internal/accountusage/service_test.go @@ -3,8 +3,10 @@ package accountusage import ( "context" "encoding/json" + "errors" "fmt" "path/filepath" + "strings" "testing" "time" @@ -23,6 +25,23 @@ type fakeRunner struct { err error } +type retryRunner struct { + calls int + failures int + snapshot codexusage.Snapshot + auth json.RawMessage + inputs []json.RawMessage +} + +func (runner *retryRunner) Query(_ context.Context, auth json.RawMessage) (codexusage.Snapshot, json.RawMessage, error) { + runner.calls++ + runner.inputs = append(runner.inputs, append(json.RawMessage(nil), auth...)) + if runner.calls <= runner.failures { + return codexusage.Snapshot{}, runner.auth, errors.New("temporary connection failure") + } + return runner.snapshot, runner.auth, nil +} + func (runner fakeRunner) Query(context.Context, json.RawMessage) (codexusage.Snapshot, json.RawMessage, error) { return runner.snapshot, runner.auth, runner.err } @@ -57,6 +76,67 @@ func TestRefreshCachesUsageAndAdoptsRotatedCredentials(t *testing.T) { } } +func TestRefreshRetriesAndClearsPersistedFailure(t *testing.T) { + service, _, profile := testService(t, false) + runner := &retryRunner{ + failures: 2, + snapshot: codexusage.Snapshot{FetchedAt: time.Now().UTC(), PlanType: "pro"}, + auth: authBytes("account-a", "refresh-new", "2026-08-20T01:00:00Z"), + } + service.Runner = runner + service.retryBaseDelay = time.Nanosecond + results, err := service.Refresh(context.Background(), []string{profile.ID}) + if err != nil { + t.Fatal(err) + } + if runner.calls != 3 || results[profile.ID].Attempts != 3 || results[profile.ID].Error != "" { + t.Fatalf("unexpected retry result: calls=%d result=%#v", runner.calls, results[profile.ID]) + } + retriedWith, err := authschema.Parse(runner.inputs[1]) + if err != nil || retriedWith.Tokens.RefreshToken != "refresh-new" { + t.Fatalf("retry did not use rotated credentials: %#v, %v", retriedWith.Tokens, err) + } + cache, err := service.Cached() + if err != nil { + t.Fatal(err) + } + if _, ok := cache.Failures[profile.ID]; ok { + t.Fatalf("successful retry retained failure: %#v", cache.Failures[profile.ID]) + } +} + +func TestRefreshPersistsFailureDetails(t *testing.T) { + service, manager, profile := testService(t, false) + runner := &retryRunner{failures: 3, auth: authBytes("account-a", "refresh-new", "2026-08-20T01:00:00Z")} + service.Runner = runner + service.retryBaseDelay = time.Nanosecond + results, err := service.Refresh(context.Background(), []string{profile.ID}) + if err != nil { + t.Fatal(err) + } + result := results[profile.ID] + if result.Attempts != 3 || !strings.Contains(result.Error, "temporary connection failure") { + t.Fatalf("unexpected failure result: %#v", result) + } + cache, err := service.Cached() + if err != nil { + t.Fatal(err) + } + failure := cache.Failures[profile.ID] + if failure.Attempts != 3 || failure.AttemptedAt.IsZero() || failure.Message != result.Error { + t.Fatalf("failure was not persisted: %#v", failure) + } + data, err := manager.Load() + if err != nil { + t.Fatal(err) + } + saved, _ := data.Find(profile.ID) + document, _ := authschema.Parse(saved.Auth) + if document.Tokens.RefreshToken != "refresh-new" { + t.Fatal("rotated credentials from failed queries were not preserved") + } +} + func TestRefreshProjectsRotatedCredentialsForActiveProfile(t *testing.T) { service, manager, profile := testService(t, true) service.Runner = fakeRunner{ diff --git a/internal/app/app.go b/internal/app/app.go index 86cb000..a323e1c 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -27,6 +27,7 @@ import ( "github.com/SilkageNet/codex-switch/internal/secretstore" appstate "github.com/SilkageNet/codex-switch/internal/state" "github.com/SilkageNet/codex-switch/internal/switcher" + "github.com/SilkageNet/codex-switch/internal/usagecache" "github.com/SilkageNet/codex-switch/internal/vault" "github.com/spf13/cobra" "golang.org/x/term" @@ -72,8 +73,12 @@ type usageView struct { TokenUsage *codexusage.TokenUsage `json:"tokenUsage,omitempty"` Partial []string `json:"partial,omitempty"` Error string `json:"error,omitempty"` + ErrorAt time.Time `json:"errorAt,omitempty"` + Attempts int `json:"attempts,omitempty"` } +const usageFreshness = 5 * time.Minute + func NewCommand(version string) *cobra.Command { options := &Options{Version: version, Input: os.Stdin, Output: os.Stdout, Error: os.Stderr} root := &cobra.Command{ @@ -289,7 +294,7 @@ func newAccountListCommand(options *Options) *cobra.Command { views := make([]accountView, 0, len(data.Profiles)) for _, profile := range data.Profiles { view := toView(profile, profile.ID == activeProfileID) - view.Usage = usageFromCache(cache.Profiles, profile.ID, refreshErrors[profile.ID], time.Now()) + view.Usage = usageFromCache(cache, profile.ID, refreshErrors[profile.ID], time.Now()) views = append(views, view) } if options.JSON { @@ -300,7 +305,7 @@ func newAccountListCommand(options *Options) *cobra.Command { return nil } writer := tabwriter.NewWriter(options.Output, 0, 4, 2, ' ', 0) - _, _ = fmt.Fprintln(writer, "\tALIAS\tIDENTITY\tPLAN\tLIMITS\tRESETS\tTOKENS\tUPDATED") + _, _ = fmt.Fprintln(writer, "\tALIAS\tIDENTITY\tPLAN\tLIMITS\tRESETS\tTOKENS\tUPDATED\tERROR") for _, view := range views { marker := " " if view.Active { @@ -311,11 +316,16 @@ func newAccountListCommand(options *Options) *cobra.Command { identity = view.AccountID } plan, limits, resets, tokens, updated := summarizeUsage(view.Usage, time.Now()) - _, _ = fmt.Fprintf(writer, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n", marker, view.Alias, identity, plan, limits, resets, tokens, updated) + _, _ = fmt.Fprintf(writer, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n", marker, view.Alias, identity, plan, limits, resets, tokens, updated, compactUsageError(view.Usage)) } if err := writer.Flush(); err != nil { return err } + for _, view := range views { + if view.Usage != nil && view.Usage.Error != "" { + _, _ = fmt.Fprintf(options.Output, "Refresh error for %s: %s\n", view.Alias, view.Usage.Error) + } + } if notice := observationNotice(observation); notice != "" { _, _ = fmt.Fprintln(options.Output, notice) } @@ -411,7 +421,7 @@ func newAccountUsageCommand(options *Options) *cobra.Command { views := make([]accountView, 0, len(profiles)) for _, profile := range profiles { view := toView(profile, profile.ID == activeProfileID) - view.Usage = usageFromCache(cache.Profiles, profile.ID, refreshErrors[profile.ID], time.Now()) + view.Usage = usageFromCache(cache, profile.ID, refreshErrors[profile.ID], time.Now()) if len(profiles) == 1 && view.Usage.Status == "unavailable" { if view.Usage.Error != "" { return errors.New(view.Usage.Error) @@ -1096,10 +1106,17 @@ func formatView(view accountView) string { return fmt.Sprintf("%s (%s)", view.Alias, identity) } -func usageFromCache(cache map[string]codexusage.Snapshot, profileID, queryError string, now time.Time) *usageView { - snapshot, ok := cache[profileID] +func usageFromCache(cache usagecache.Cache, profileID, queryError string, now time.Time) *usageView { + failure := cache.Failures[profileID] + if queryError == "" { + queryError = failure.Message + } + if queryError != "" && failure.AttemptedAt.IsZero() { + failure.AttemptedAt = now + } + snapshot, ok := cache.Profiles[profileID] if !ok { - return &usageView{Status: "unavailable", Error: queryError} + return &usageView{Status: "unavailable", Error: queryError, ErrorAt: failure.AttemptedAt, Attempts: failure.Attempts} } return &usageView{ Status: usageStatus(snapshot.FetchedAt, now), @@ -1109,16 +1126,31 @@ func usageFromCache(cache map[string]codexusage.Snapshot, profileID, queryError TokenUsage: snapshot.TokenUsage, Partial: snapshot.Partial, Error: queryError, + ErrorAt: failure.AttemptedAt, + Attempts: failure.Attempts, } } func usageStatus(fetchedAt, now time.Time) string { - if fetchedAt.IsZero() || now.Sub(fetchedAt) > time.Minute { + if fetchedAt.IsZero() || now.Sub(fetchedAt) > usageFreshness { return "stale" } return "fresh" } +func compactUsageError(usage *usageView) string { + if usage == nil || usage.Error == "" { + return "-" + } + value := strings.Join(strings.Fields(usage.Error), " ") + const maxRunes = 56 + runes := []rune(value) + if len(runes) > maxRunes { + value = string(runes[:maxRunes-1]) + "…" + } + return value +} + func summarizeUsage(usage *usageView, now time.Time) (string, string, string, string, string) { if usage == nil || usage.Status == "unavailable" { return "-", "unavailable", "-", "-", "-" @@ -1236,7 +1268,14 @@ func formatUsage(view accountView, now time.Time) string { output.WriteString(strings.Join(usage.Partial, ", ")) } if usage.Error != "" { - output.WriteString("\nWarning: ") + output.WriteString("\nRefresh error") + if usage.Attempts > 0 { + fmt.Fprintf(&output, " after %d attempt(s)", usage.Attempts) + } + if !usage.ErrorAt.IsZero() { + fmt.Fprintf(&output, " at %s", formatLocalTime(usage.ErrorAt)) + } + output.WriteString(": ") output.WriteString(usage.Error) } return output.String() diff --git a/internal/app/usage_test.go b/internal/app/usage_test.go index 6070a22..7741317 100644 --- a/internal/app/usage_test.go +++ b/internal/app/usage_test.go @@ -7,6 +7,7 @@ import ( "github.com/SilkageNet/codex-switch/internal/codexusage" "github.com/SilkageNet/codex-switch/internal/switcher" + "github.com/SilkageNet/codex-switch/internal/usagecache" ) func TestSummarizeUsage(t *testing.T) { @@ -38,6 +39,24 @@ func TestSummarizeUsage(t *testing.T) { } } +func TestUsageFromCacheRetainsLastRefreshFailure(t *testing.T) { + now := time.Date(2026, 8, 21, 12, 0, 0, 0, time.UTC) + cache := usagecache.Cache{ + Profiles: map[string]codexusage.Snapshot{"profile-a": {FetchedAt: now.Add(-2 * time.Minute), PlanType: "pro"}}, + Failures: map[string]usagecache.Failure{"profile-a": {AttemptedAt: now.Add(-time.Minute), Message: "temporary connection failure", Attempts: 3}}, + } + view := usageFromCache(cache, "profile-a", "", now) + if view.Status != "fresh" || view.Error != "temporary connection failure" || view.Attempts != 3 { + t.Fatalf("unexpected cached usage view: %#v", view) + } + if got := compactUsageError(view); got != "temporary connection failure" { + t.Fatalf("compact error = %q", got) + } + if usageStatus(now.Add(-6*time.Minute), now) != "stale" { + t.Fatal("snapshot older than five minutes was not stale") + } +} + func TestObservationMessagesExplainExternalLogin(t *testing.T) { observation := switcher.Observation{ State: switcher.AccountStateExternalLoginWithRefresh, diff --git a/internal/codexusage/runner.go b/internal/codexusage/runner.go index 2ceb638..7616de8 100644 --- a/internal/codexusage/runner.go +++ b/internal/codexusage/runner.go @@ -125,23 +125,37 @@ func (runner Runner) Query(ctx context.Context, auth json.RawMessage) (Snapshot, if err := encoder.Encode(rpcRequest{Method: "initialized", Params: map[string]any{}}); err != nil { return Snapshot{}, nil, fmt.Errorf("send initialized notification: %w", err) } - requests := []struct { - id int - method string - params any - }{ - {id: 1, method: "account/read", params: map[string]bool{"refreshToken": false}}, - {id: 2, method: "account/rateLimits/read"}, - {id: 3, method: "account/usage/read"}, - } - for _, request := range requests { - if err := sendRequest(encoder, request.id, request.method, request.params); err != nil { - return Snapshot{}, nil, err - } + if err := sendRequest(encoder, 1, "account/read", map[string]bool{"refreshToken": false}); err != nil { + return Snapshot{}, nil, err } - responses, err := readResponses(scanner, 1, 2, 3) + accountResponse, err := readResponse(scanner, 1) if err != nil { - return Snapshot{}, nil, protocolError("read account usage", err, stderr.String()) + return Snapshot{}, nil, protocolError("read account", err, stderr.String()) + } + if accountResponse.Error != nil { + return Snapshot{}, nil, responseError("account/read", accountResponse.Error) + } + var account AccountResponse + if err := json.Unmarshal(accountResponse.Result, &account); err != nil { + return Snapshot{}, nil, fmt.Errorf("decode account/read response: %w", err) + } + if account.Account == nil || account.Account.Type != "chatgpt" { + return Snapshot{}, nil, errors.New("saved profile is not recognized as a ChatGPT account by Codex") + } + + if err := sendRequest(encoder, 2, "account/rateLimits/read", nil); err != nil { + return Snapshot{}, nil, err + } + rateLimitsResponse, err := readResponse(scanner, 2) + if err != nil { + return Snapshot{}, nil, protocolError("read account rate limits", err, stderr.String()) + } + if err := sendRequest(encoder, 3, "account/usage/read", nil); err != nil { + return Snapshot{}, nil, err + } + tokenUsageResponse, err := readResponse(scanner, 3) + if err != nil { + return Snapshot{}, nil, protocolError("read account token usage", err, stderr.String()) } _ = stdin.Close() if err := command.Wait(); err != nil { @@ -152,51 +166,42 @@ func (runner Runner) Query(ctx context.Context, auth json.RawMessage) (Snapshot, return Snapshot{}, nil, protocolError("Codex app server exited", err, stderr.String()) } waited = true - - accountResponse := responses[1] - if accountResponse.Error != nil { - return Snapshot{}, nil, responseError("account/read", accountResponse.Error) + updated, err := atomicfile.ReadLimited(filepath.Join(temporaryHome, "auth.json"), 2<<20) + if err != nil { + return Snapshot{}, nil, fmt.Errorf("read refreshed account credentials: %w", err) } - var account AccountResponse - if err := json.Unmarshal(accountResponse.Result, &account); err != nil { - return Snapshot{}, nil, fmt.Errorf("decode account/read response: %w", err) + updatedDocument, err := authschema.Parse(updated) + if err != nil { + return Snapshot{}, nil, fmt.Errorf("validate refreshed account credentials: %w", err) } - if account.Account == nil || account.Account.Type != "chatgpt" { - return Snapshot{}, nil, errors.New("saved profile is not recognized as a ChatGPT account by Codex") + if updatedDocument.Tokens.AccountID != document.Tokens.AccountID { + return Snapshot{}, nil, errors.New("codex returned credentials for a different account") } + snapshot := Snapshot{FetchedAt: time.Now().UTC(), PlanType: account.Account.PlanType} - if response := responses[2]; response.Error != nil { + if rateLimitsResponse.Error != nil { snapshot.Partial = append(snapshot.Partial, "rateLimits") } else { var value RateLimits - if err := json.Unmarshal(response.Result, &value); err != nil { + if err := json.Unmarshal(rateLimitsResponse.Result, &value); err != nil { return Snapshot{}, nil, fmt.Errorf("decode account/rateLimits/read response: %w", err) } snapshot.RateLimits = &value } - if response := responses[3]; response.Error != nil { + if tokenUsageResponse.Error != nil { snapshot.Partial = append(snapshot.Partial, "tokenUsage") } else { var value TokenUsage - if err := json.Unmarshal(response.Result, &value); err != nil { + if err := json.Unmarshal(tokenUsageResponse.Result, &value); err != nil { return Snapshot{}, nil, fmt.Errorf("decode account/usage/read response: %w", err) } snapshot.TokenUsage = &value } if snapshot.RateLimits == nil && snapshot.TokenUsage == nil { - return Snapshot{}, nil, errors.New("this Codex version does not provide account usage methods; update Codex and retry") - } - - updated, err := atomicfile.ReadLimited(filepath.Join(temporaryHome, "auth.json"), 2<<20) - if err != nil { - return Snapshot{}, nil, fmt.Errorf("read refreshed account credentials: %w", err) - } - updatedDocument, err := authschema.Parse(updated) - if err != nil { - return Snapshot{}, nil, fmt.Errorf("validate refreshed account credentials: %w", err) - } - if updatedDocument.Tokens.AccountID != document.Tokens.AccountID { - return Snapshot{}, nil, errors.New("codex returned credentials for a different account") + return Snapshot{}, updatedDocument.Raw, fmt.Errorf("account usage requests failed: %w", errors.Join( + responseError("account/rateLimits/read", rateLimitsResponse.Error), + responseError("account/usage/read", tokenUsageResponse.Error), + )) } return snapshot, updatedDocument.Raw, nil } diff --git a/internal/codexusage/runner_test.go b/internal/codexusage/runner_test.go index 42bc951..7b3bd95 100644 --- a/internal/codexusage/runner_test.go +++ b/internal/codexusage/runner_test.go @@ -58,6 +58,22 @@ func TestRunnerAllowsOneUnsupportedUsageMethod(t *testing.T) { } } +func TestRunnerPreservesBothUsageMethodErrors(t *testing.T) { + runner := helperRunnerMode(t, false, true) + _, updated, err := runner.Query(context.Background(), testAuth("account-a", "refresh-old", "2026-08-20T00:00:00Z")) + if err == nil { + t.Fatal("expected usage query error") + } + for _, expected := range []string{"account/rateLimits/read failed (-32001): overloaded", "account/usage/read failed (-32601): method not found"} { + if !strings.Contains(err.Error(), expected) { + t.Fatalf("error %q does not contain %q", err, expected) + } + } + if !strings.Contains(string(updated), "refresh-new") { + t.Fatalf("rotated credentials were lost on usage failure: %s", updated) + } +} + func TestProtocolErrorRedactsChildDiagnostics(t *testing.T) { err := protocolError("query", context.DeadlineExceeded, "secret child output") if strings.Contains(err.Error(), "secret child output") || !strings.Contains(err.Error(), "redacted") { @@ -74,6 +90,10 @@ func TestRunnerReturnsBinaryDiscoveryError(t *testing.T) { } func helperRunner(t *testing.T, partial bool) Runner { + return helperRunnerMode(t, partial, false) +} + +func helperRunnerMode(t *testing.T, partial, allError bool) Runner { t.Helper() return Runner{ Binary: os.Args[0], @@ -84,6 +104,9 @@ func helperRunner(t *testing.T, partial bool) Runner { if partial { command.Env = append(command.Env, "CODEX_USAGE_PARTIAL=1") } + if allError { + command.Env = append(command.Env, "CODEX_USAGE_ALL_ERROR=1") + } return command }, } @@ -113,6 +136,10 @@ func TestCodexUsageHelperProcess(t *testing.T) { "requiresOpenaiAuth": true, }}) case "account/rateLimits/read": + if os.Getenv("CODEX_USAGE_ALL_ERROR") == "1" { + _ = encoder.Encode(map[string]any{"id": json.RawMessage(id), "error": map[string]any{"code": -32001, "message": "overloaded"}}) + continue + } _ = encoder.Encode(map[string]any{"id": json.RawMessage(id), "result": map[string]any{ "rateLimits": map[string]any{"planType": "pro", "primary": map[string]any{"usedPercent": 21, "windowDurationMins": 300}}, "rateLimitsByLimitId": map[string]any{"codex": map[string]any{"planType": "pro", "primary": map[string]any{"usedPercent": 21, "windowDurationMins": 300}}}, @@ -130,7 +157,7 @@ func TestCodexUsageHelperProcess(t *testing.T) { }, }}) case "account/usage/read": - if os.Getenv("CODEX_USAGE_PARTIAL") == "1" { + if os.Getenv("CODEX_USAGE_PARTIAL") == "1" || os.Getenv("CODEX_USAGE_ALL_ERROR") == "1" { _ = encoder.Encode(map[string]any{"id": json.RawMessage(id), "error": map[string]any{"code": -32601, "message": "method not found"}}) } else { lifetime := int64(1234567) diff --git a/internal/usagecache/cache.go b/internal/usagecache/cache.go index c00519b..855095a 100644 --- a/internal/usagecache/cache.go +++ b/internal/usagecache/cache.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "os" + "time" "github.com/SilkageNet/codex-switch/internal/atomicfile" "github.com/SilkageNet/codex-switch/internal/codexusage" @@ -13,12 +14,19 @@ import ( type Cache struct { Version int `json:"version"` Profiles map[string]codexusage.Snapshot `json:"profiles"` + Failures map[string]Failure `json:"failures,omitempty"` +} + +type Failure struct { + AttemptedAt time.Time `json:"attemptedAt"` + Message string `json:"message"` + Attempts int `json:"attempts,omitempty"` } func Load(path string) (Cache, error) { data, err := atomicfile.ReadLimited(path, 16<<20) if errors.Is(err, os.ErrNotExist) { - return Cache{Version: 1, Profiles: map[string]codexusage.Snapshot{}}, nil + return Cache{Version: 1, Profiles: map[string]codexusage.Snapshot{}, Failures: map[string]Failure{}}, nil } if err != nil { return Cache{}, err @@ -33,6 +41,9 @@ func Load(path string) (Cache, error) { if cache.Profiles == nil { cache.Profiles = map[string]codexusage.Snapshot{} } + if cache.Failures == nil { + cache.Failures = map[string]Failure{} + } return cache, nil } @@ -41,6 +52,9 @@ func Save(path string, cache Cache) error { if cache.Profiles == nil { cache.Profiles = map[string]codexusage.Snapshot{} } + if cache.Failures == nil { + cache.Failures = map[string]Failure{} + } data, err := json.MarshalIndent(cache, "", " ") if err != nil { return err diff --git a/internal/usagecache/cache_test.go b/internal/usagecache/cache_test.go index 01b7724..a543e5e 100644 --- a/internal/usagecache/cache_test.go +++ b/internal/usagecache/cache_test.go @@ -16,7 +16,7 @@ func TestLoadMissingAndRoundTrip(t *testing.T) { if err != nil { t.Fatal(err) } - if cache.Version != 1 || cache.Profiles == nil { + if cache.Version != 1 || cache.Profiles == nil || cache.Failures == nil { t.Fatalf("unexpected empty cache: %#v", cache) } fetched := time.Date(2026, 8, 20, 1, 2, 3, 0, time.UTC) @@ -29,6 +29,7 @@ func TestLoadMissingAndRoundTrip(t *testing.T) { Credits: []codexusage.RateLimitResetCredit{{ID: "reset-1", Status: "available", ExpiresAt: &expires}}, }}, } + cache.Failures["profile-b"] = Failure{AttemptedAt: fetched, Message: "network unavailable", Attempts: 3} if err := Save(path, cache); err != nil { t.Fatal(err) } @@ -39,6 +40,9 @@ func TestLoadMissingAndRoundTrip(t *testing.T) { if loaded.Profiles["profile-a"].PlanType != "pro" || !loaded.Profiles["profile-a"].FetchedAt.Equal(fetched) { t.Fatalf("unexpected cache round trip: %#v", loaded) } + if failure := loaded.Failures["profile-b"]; failure.Message != "network unavailable" || failure.Attempts != 3 || !failure.AttemptedAt.Equal(fetched) { + t.Fatalf("failure was not preserved: %#v", failure) + } resetCredits := loaded.Profiles["profile-a"].RateLimits.RateLimitResetCredits if resetCredits == nil || resetCredits.AvailableCount != 2 || len(resetCredits.Credits) != 1 || resetCredits.Credits[0].ExpiresAt == nil || *resetCredits.Credits[0].ExpiresAt != expires { t.Fatalf("reset credits were not preserved: %#v", resetCredits)