From 28dd7a4fe2866b9f88dc10be05c67140c345a2ce Mon Sep 17 00:00:00 2001 From: Silkage Date: Fri, 21 Aug 2026 11:25:48 +0800 Subject: [PATCH] feat: query account usage without switching --- README.md | 37 ++- docs/architecture.md | 28 ++ docs/compatibility.md | 21 +- docs/security.md | 7 + docs/troubleshooting.md | 25 ++ internal/accountusage/service.go | 274 +++++++++++++++++ internal/accountusage/service_test.go | 160 ++++++++++ internal/app/app.go | 423 +++++++++++++++++++++++++- internal/app/usage_test.go | 58 ++++ internal/codexusage/runner.go | 287 +++++++++++++++++ internal/codexusage/runner_test.go | 124 ++++++++ internal/codexusage/types.go | 104 +++++++ internal/config/paths.go | 18 +- internal/usagecache/cache.go | 49 +++ internal/usagecache/cache_test.go | 43 +++ 15 files changed, 1631 insertions(+), 27 deletions(-) create mode 100644 internal/accountusage/service.go create mode 100644 internal/accountusage/service_test.go create mode 100644 internal/app/usage_test.go create mode 100644 internal/codexusage/runner.go create mode 100644 internal/codexusage/runner_test.go create mode 100644 internal/codexusage/types.go create mode 100644 internal/usagecache/cache.go create mode 100644 internal/usagecache/cache_test.go diff --git a/README.md b/README.md index db5e19d..6b68d00 100644 --- a/README.md +++ b/README.md @@ -55,12 +55,39 @@ codex-switch account add work --device-auth # Inspect and switch. Close Codex before switching, then restart it. codex-switch account list +codex-switch account usage work codex-switch use work codex-switch current ``` Use `codex-switch doctor` before reporting a problem. Machine-readable output is -available on status and list commands with `--json`. +available on commands with `--json`. + +## Usage without switching + +`codex-switch` can inspect every saved account through the official Codex App +Server without making that account active: + +```bash +# Query the active managed account now. +codex-switch account usage + +# Query one saved account, or all accounts concurrently. +codex-switch account usage work +codex-switch account usage --all + +# Refresh all rows in the compact account table. +codex-switch account list --refresh + +# Work offline with the last successful snapshots. +codex-switch account list --cached +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 +not switch `$CODEX_HOME/auth.json`, sessions, plugins, or UI state. The cache +contains usage numbers and public account metadata only, never tokens. ## Commands @@ -75,7 +102,8 @@ codex-switch select codex-switch account add codex-switch account import-current -codex-switch account list +codex-switch account list [--refresh|--cached] +codex-switch account usage [alias] [--all] [--cached] codex-switch account show codex-switch account rename codex-switch account reauth @@ -94,6 +122,11 @@ Normal account switches modify only: - `$CODEX_HOME/auth.json` - `codex-switch`'s own state and encrypted vault +An account-usage query may persist an officially refreshed credential generation +back to the encrypted vault. If the queried profile is active, the same +generation is safely reconciled into `$CODEX_HOME/auth.json`; no account selection +or Codex-owned state changes. + Initialization may make a one-time, backed-up change to `$CODEX_HOME/config.toml` to set `cli_auth_credentials_store = "file"`. `codex-switch` does not rewrite session history or Codex configuration during a diff --git a/docs/architecture.md b/docs/architecture.md index 5234e78..0c09bdd 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -20,6 +20,12 @@ - `vault` encrypts all saved account profiles with XChaCha20-Poly1305. - `switcher` reconciles a live Codex refresh generation, prepares a journal, performs compare-before-replace, and records the selected profile. +- `codexusage` runs the official Codex App Server in an isolated temporary + `CODEX_HOME` and reads the stable account, rate-limit, and token-usage methods. +- `accountusage` queries up to four profiles concurrently, reconciles credential + refresh generations, and coordinates with switching through the same lock. +- `usagecache` stores credential-free successful snapshots separately from the + encrypted vault. - `atomicfile` publishes complete files and refuses symlink destinations. - `doctor` reports only redacted, non-secret local facts. @@ -43,6 +49,25 @@ The journal contains only profile IDs, hashes, and timestamps. If the process stops after replacement but before state persistence, recovery compares the live file with both hashes and completes the state transition. +## Isolated usage query + +```text +acquire shared operation lock + -> decrypt and validate selected profile(s) + -> 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 + -> stop the server and delete the temporary home + -> reconcile any newer credential generation + -> 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. + ## Data locations `CODEX_HOME` resolution: @@ -59,3 +84,6 @@ file with both hashes and completes the state transition. Only tests and advanced portable installations should normally override these paths. + +The usage cache is `usage-cache.v1.json` inside the resolved `codex-switch` data +directory. It contains no authentication documents. diff --git a/docs/compatibility.md b/docs/compatibility.md index 00059d4..8e7156b 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -27,10 +27,21 @@ The command creates a timestamped backup before making a surgical top-level ## Codex releases -Development began against Codex CLI `0.148.0-alpha.15`. The project does not use -that version's private OAuth endpoints. Login is delegated to the installed -official CLI, reducing the compatibility surface to its documented cached -authentication shape. +Development began against Codex CLI `0.148.0-alpha.15`; isolated account-usage +queries were validated with `0.148.0-alpha.21`. The project does not use private +OAuth or usage endpoints. Login is delegated to the installed official CLI, and +usage is read through the documented stable Codex App Server protocol. + +Usage querying initializes `codex app-server` and calls: + +- `account/read` +- `account/rateLimits/read` +- `account/usage/read` + +If one usage method is unavailable, the other is still cached and marked +partial. If both are unavailable, update the installed Codex client. These +methods require a ChatGPT/Codex service login; API-key-only and Amazon Bedrock +profiles are not supported by `codex-switch`. On an unknown or malformed schema, `codex-switch` stops before overwriting the live file. Add a redacted fixture and a versioned adapter before broadening the @@ -39,5 +50,7 @@ accepted shape. ## Upstream references - OpenAI authentication documentation: https://developers.openai.com/codex/auth +- OpenAI Codex App Server documentation: + https://learn.chatgpt.com/docs/app-server - CC Switch managed Codex OAuth implementation: https://github.com/farion1231/cc-switch/tree/v3.20.0 diff --git a/docs/security.md b/docs/security.md index 4057bff..9db9294 100644 --- a/docs/security.md +++ b/docs/security.md @@ -32,6 +32,13 @@ focuses on: - Commands never expose a token retrieval operation. - JSON output is intentionally based on dedicated public view types. +- Usage queries create per-profile temporary homes with mode `0700` and + credential files with mode `0600`, then remove them after the App Server exits. +- The usage cache contains rate limits, aggregate token statistics, timestamps, + and public account metadata only. It 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. - Real credentials are forbidden in tests and fixtures. - The Linux desktop implementation fails closed when Secret Service is absent. - Portable backups require a passphrase of at least 12 characters and use diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 92e1e50..1a0c6b5 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -47,6 +47,31 @@ codex-switch account reauth The ambiguity is intentionally not resolved by guessing. +## Account usage is unavailable + +Confirm the official Codex executable is installed and current: + +```bash +codex --version +codex-switch account usage +``` + +Usage queries require network access and a saved ChatGPT login. They do not work +for API-key-only or Amazon Bedrock authentication. A missing method on an older +Codex build is reported as partial when the other method still works; update +Codex if both usage methods are unavailable. + +For an offline or temporarily failing service, inspect the last successful +snapshot without making a request: + +```bash +codex-switch account list --cached +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. + ## Interrupted switch journal Run `codex-switch status`. Recovery compares the current `auth.json` hash with diff --git a/internal/accountusage/service.go b/internal/accountusage/service.go new file mode 100644 index 0000000..1bf6488 --- /dev/null +++ b/internal/accountusage/service.go @@ -0,0 +1,274 @@ +package accountusage + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "sync" + "time" + + "github.com/SilkageNet/codex-switch/internal/atomicfile" + "github.com/SilkageNet/codex-switch/internal/authschema" + "github.com/SilkageNet/codex-switch/internal/codexhome" + "github.com/SilkageNet/codex-switch/internal/codexusage" + appconfig "github.com/SilkageNet/codex-switch/internal/config" + "github.com/SilkageNet/codex-switch/internal/filelock" + appstate "github.com/SilkageNet/codex-switch/internal/state" + "github.com/SilkageNet/codex-switch/internal/usagecache" + "github.com/SilkageNet/codex-switch/internal/vault" +) + +const queryTimeout = 20 * time.Second + +type QueryRunner interface { + Query(context.Context, json.RawMessage) (codexusage.Snapshot, json.RawMessage, error) +} + +type Service struct { + Home codexhome.Home + Paths appconfig.Paths + Vault *vault.Manager + Runner QueryRunner +} + +type Result struct { + ProfileID string `json:"profileId"` + Snapshot codexusage.Snapshot `json:"snapshot,omitempty"` + Error string `json:"error,omitempty"` +} + +type queryResult struct { + profileID string + snapshot codexusage.Snapshot + auth json.RawMessage + err error +} + +func (service Service) Cached() (usagecache.Cache, error) { + return usagecache.Load(service.Paths.UsageCache) +} + +func (service Service) Refresh(ctx context.Context, profileIDs []string) (map[string]Result, error) { + if service.Runner == nil { + return nil, errors.New("account usage runner is not configured") + } + lock, err := filelock.Acquire(filepath.Join(service.Home.Path, ".codex-switch.lock")) + if err != nil { + return nil, err + } + defer func() { _ = lock.Close() }() + + data, err := service.Vault.Load() + if err != nil { + return nil, err + } + selected := make([]vault.Profile, 0, len(profileIDs)) + seen := make(map[string]bool, len(profileIDs)) + for _, id := range profileIDs { + profile, findErr := data.Find(id) + if findErr != nil { + return nil, findErr + } + if !seen[profile.ID] { + selected = append(selected, *profile) + seen[profile.ID] = true + } + } + + queried := make(chan queryResult, len(selected)) + semaphore := make(chan struct{}, 4) + var group sync.WaitGroup + 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} + }() + } + group.Wait() + close(queried) + + results := make(map[string]Result, len(selected)) + candidates := make(map[string]json.RawMessage, len(selected)) + cache, err := usagecache.Load(service.Paths.UsageCache) + if err != nil { + return nil, err + } + for result := range queried { + entry := Result{ProfileID: result.profileID} + if result.err != nil { + entry.Error = result.err.Error() + } else { + entry.Snapshot = result.snapshot + candidates[result.profileID] = result.auth + } + results[result.profileID] = entry + } + + data, err = service.Vault.Load() + if err != nil { + return nil, err + } + state, stateErr := appstate.Load(service.Paths.State) + if stateErr != nil && !errors.Is(stateErr, os.ErrNotExist) { + return nil, stateErr + } + vaultChanged := false + for profileID, candidateRaw := range candidates { + profile, findErr := data.Find(profileID) + if findErr != nil { + entry := results[profileID] + entry.Error = appendError(entry.Error, "profile was removed while usage was being queried") + results[profileID] = entry + delete(cache.Profiles, profileID) + continue + } + candidate, parseErr := authschema.Parse(candidateRaw) + if parseErr != nil { + entry := results[profileID] + entry.Error = appendError(entry.Error, "refreshed credentials are invalid") + results[profileID] = entry + continue + } + if candidate.Tokens.AccountID != profile.AccountID || (profile.WorkspaceID != "" && candidate.WorkspaceID != "" && candidate.WorkspaceID != profile.WorkspaceID) { + entry := results[profileID] + entry.Error = appendError(entry.Error, "refreshed credentials do not match the saved account") + results[profileID] = entry + continue + } + cache.Profiles[profileID] = results[profileID].Snapshot + saved, parseErr := authschema.Parse(profile.Auth) + if parseErr != nil { + entry := results[profileID] + entry.Error = appendError(entry.Error, "saved credentials became invalid") + results[profileID] = entry + continue + } + decision, compareErr := authschema.CompareGeneration(saved, candidate) + if compareErr != nil { + entry := results[profileID] + entry.Error = appendError(entry.Error, credentialError(compareErr)) + results[profileID] = entry + continue + } + if decision != authschema.GenerationAdoptLive { + continue + } + profile.Auth = append(json.RawMessage(nil), candidate.Raw...) + profile.Email = candidate.Email + profile.WorkspaceID = candidate.WorkspaceID + if refreshed, ok := candidate.GenerationTime(); ok { + profile.TokenUpdatedAt = refreshed + } + vaultChanged = true + + if state.ActiveProfileID == profile.ID { + changed, syncErr := service.reconcileActive(profile, candidate, &state) + if syncErr != nil { + entry := results[profileID] + entry.Error = appendError(entry.Error, syncErr.Error()) + results[profileID] = entry + } else if changed { + vaultChanged = true + } + } + } + if vaultChanged { + if err := service.Vault.Save(data); err != nil { + return nil, err + } + } + validProfiles := make(map[string]bool, len(data.Profiles)) + for _, profile := range data.Profiles { + validProfiles[profile.ID] = true + } + for profileID := range cache.Profiles { + if !validProfiles[profileID] { + delete(cache.Profiles, profileID) + } + } + if err := usagecache.Save(service.Paths.UsageCache, cache); err != nil { + return nil, err + } + return results, nil +} + +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) { + return false, nil + } + if err != nil { + return false, fmt.Errorf("read active credentials after refresh: %w", err) + } + liveHash := atomicfile.Hash(liveRaw) + live, err := authschema.Parse(liveRaw) + if err != nil { + return false, fmt.Errorf("validate active credentials after refresh: %w", err) + } + if live.Tokens.AccountID != candidate.Tokens.AccountID { + return false, errors.New("active account changed while usage was being queried; refreshed credentials were kept only in the vault") + } + decision, err := authschema.CompareGeneration(candidate, live) + if err != nil { + return false, errors.New(credentialError(err)) + } + switch decision { + case authschema.GenerationAdoptLive: + profile.Auth = append(json.RawMessage(nil), live.Raw...) + profile.Email = live.Email + profile.WorkspaceID = live.WorkspaceID + if refreshed, ok := live.GenerationTime(); ok { + profile.TokenUpdatedAt = refreshed + } + return true, nil + case authschema.GenerationUseSaved: + currentHash, hashErr := service.Home.AuthHash() + if hashErr != nil { + return false, hashErr + } + if currentHash != liveHash { + return false, errors.New("active credentials changed while usage was being queried; no active file was replaced") + } + if err := service.Home.WriteAuth(candidate.Raw); err != nil { + return false, fmt.Errorf("publish refreshed active credentials: %w", err) + } + publishedHash, err := service.Home.AuthHash() + if err != nil { + return false, err + } + if err := appstate.Save(service.Paths.State, appstate.State{ActiveProfileID: profile.ID, AuthHash: publishedHash}); err != nil { + return false, fmt.Errorf("record refreshed active credentials: %w", err) + } + *state = appstate.State{Version: 1, ActiveProfileID: profile.ID, AuthHash: publishedHash} + } + return false, nil +} + +func credentialError(err error) string { + if errors.Is(err, authschema.ErrAmbiguousGeneration) { + return "refreshed Token generations are ambiguous; reauthenticate this profile" + } + return err.Error() +} + +func appendError(existing, next string) string { + if existing == "" { + return next + } + return existing + "; " + next +} diff --git a/internal/accountusage/service_test.go b/internal/accountusage/service_test.go new file mode 100644 index 0000000..1b060e7 --- /dev/null +++ b/internal/accountusage/service_test.go @@ -0,0 +1,160 @@ +package accountusage + +import ( + "context" + "encoding/json" + "fmt" + "path/filepath" + "testing" + "time" + + "github.com/SilkageNet/codex-switch/internal/authschema" + "github.com/SilkageNet/codex-switch/internal/codexhome" + "github.com/SilkageNet/codex-switch/internal/codexusage" + appconfig "github.com/SilkageNet/codex-switch/internal/config" + "github.com/SilkageNet/codex-switch/internal/secretstore" + appstate "github.com/SilkageNet/codex-switch/internal/state" + "github.com/SilkageNet/codex-switch/internal/vault" +) + +type fakeRunner struct { + snapshot codexusage.Snapshot + auth json.RawMessage + err error +} + +func (runner fakeRunner) Query(context.Context, json.RawMessage) (codexusage.Snapshot, json.RawMessage, error) { + return runner.snapshot, runner.auth, runner.err +} + +func TestRefreshCachesUsageAndAdoptsRotatedCredentials(t *testing.T) { + service, manager, profile := testService(t, false) + updated := authBytes("account-a", "refresh-new", "2026-08-20T01:00:00Z") + lifetime := int64(1234) + service.Runner = fakeRunner{ + snapshot: codexusage.Snapshot{FetchedAt: time.Now().UTC(), PlanType: "pro", TokenUsage: &codexusage.TokenUsage{Summary: codexusage.TokenUsageSummary{LifetimeTokens: &lifetime}}}, + auth: updated, + } + results, err := service.Refresh(context.Background(), []string{profile.ID}) + if err != nil { + t.Fatal(err) + } + if results[profile.ID].Error != "" { + t.Fatalf("unexpected refresh warning: %s", results[profile.ID].Error) + } + loaded, err := manager.Load() + if err != nil { + t.Fatal(err) + } + saved, _ := loaded.Find(profile.ID) + document, _ := authschema.Parse(saved.Auth) + if document.Tokens.RefreshToken != "refresh-new" { + t.Fatal("rotated credential was not saved") + } + cache, err := service.Cached() + if err != nil || cache.Profiles[profile.ID].PlanType != "pro" { + t.Fatalf("usage was not cached: %#v, %v", cache, err) + } +} + +func TestRefreshProjectsRotatedCredentialsForActiveProfile(t *testing.T) { + service, manager, profile := testService(t, true) + service.Runner = fakeRunner{ + snapshot: codexusage.Snapshot{FetchedAt: time.Now().UTC(), PlanType: "plus"}, + auth: authBytes("account-a", "refresh-new", "2026-08-20T01:00:00Z"), + } + results, err := service.Refresh(context.Background(), []string{profile.ID}) + if err != nil { + t.Fatal(err) + } + if results[profile.ID].Error != "" { + t.Fatalf("unexpected refresh warning: %s", results[profile.ID].Error) + } + live, err := service.Home.ReadAuth() + if err != nil { + t.Fatal(err) + } + liveDocument, _ := authschema.Parse(live) + if liveDocument.Tokens.RefreshToken != "refresh-new" { + t.Fatal("active credential projection was not refreshed") + } + loaded, _ := manager.Load() + saved, _ := loaded.Find(profile.ID) + savedDocument, _ := authschema.Parse(saved.Auth) + if savedDocument.Tokens.RefreshToken != "refresh-new" { + t.Fatal("active vault credential was not refreshed") + } +} + +func TestRefreshRejectsCredentialsForDifferentAccount(t *testing.T) { + service, manager, profile := testService(t, false) + service.Runner = fakeRunner{ + snapshot: codexusage.Snapshot{FetchedAt: time.Now().UTC(), PlanType: "pro"}, + auth: authBytes("account-other", "refresh-new", "2026-08-20T01:00:00Z"), + } + results, err := service.Refresh(context.Background(), []string{profile.ID}) + if err != nil { + t.Fatal(err) + } + if results[profile.ID].Error == "" { + t.Fatal("expected credential identity warning") + } + loaded, _ := manager.Load() + saved, _ := loaded.Find(profile.ID) + document, _ := authschema.Parse(saved.Auth) + if document.Tokens.RefreshToken != "refresh-old" { + t.Fatal("mismatched credential was saved") + } + cache, cacheErr := service.Cached() + if cacheErr != nil { + t.Fatal(cacheErr) + } + if _, ok := cache.Profiles[profile.ID]; ok { + t.Fatal("usage for mismatched credentials was cached") + } +} + +func testService(t *testing.T, active bool) (Service, *vault.Manager, vault.Profile) { + t.Helper() + root := t.TempDir() + home, err := codexhome.Resolve(filepath.Join(root, "codex")) + if err != nil || home.Ensure() != nil { + t.Fatal(err) + } + paths, err := appconfig.ResolvePaths(filepath.Join(root, "switch")) + if err != nil || paths.Ensure() != nil { + t.Fatal(err) + } + manager := vault.New(paths.Vault, secretstore.NewMemoryStore()) + data, err := manager.Init() + if err != nil { + t.Fatal(err) + } + document, err := authschema.Parse(authBytes("account-a", "refresh-old", "2026-08-20T00:00:00Z")) + if err != nil { + t.Fatal(err) + } + updatedAt, _ := document.GenerationTime() + profile := vault.NewProfile("a", "test", document.Raw, "account-a", "", "a@example.com", updatedAt) + if err := data.Add(profile, false); err != nil { + t.Fatal(err) + } + if err := manager.Save(data); err != nil { + t.Fatal(err) + } + saved, _ := data.Find("a") + if active { + if err := home.WriteAuth(saved.Auth); err != nil { + t.Fatal(err) + } + hash, _ := home.AuthHash() + if err := appstate.Save(paths.State, appstate.State{ActiveProfileID: saved.ID, AuthHash: hash}); err != nil { + t.Fatal(err) + } + } + return Service{Home: home, Paths: paths, Vault: manager}, manager, *saved +} + +func authBytes(account, refresh, lastRefresh string) []byte { + return []byte(fmt.Sprintf(`{"auth_mode":"chatgpt","tokens":{"id_token":%q,"access_token":%q,"refresh_token":%q,"account_id":%q},"last_refresh":%q}`, "id-"+refresh, "access-"+refresh, refresh, account, lastRefresh)) +} diff --git a/internal/app/app.go b/internal/app/app.go index 41cb8c0..9131310 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -8,15 +8,19 @@ import ( "fmt" "io" "os" + "sort" "strconv" "strings" + "text/tabwriter" "time" + "github.com/SilkageNet/codex-switch/internal/accountusage" "github.com/SilkageNet/codex-switch/internal/atomicfile" "github.com/SilkageNet/codex-switch/internal/authschema" "github.com/SilkageNet/codex-switch/internal/cliupdate" "github.com/SilkageNet/codex-switch/internal/codexhome" "github.com/SilkageNet/codex-switch/internal/codexlogin" + "github.com/SilkageNet/codex-switch/internal/codexusage" appconfig "github.com/SilkageNet/codex-switch/internal/config" "github.com/SilkageNet/codex-switch/internal/doctor" "github.com/SilkageNet/codex-switch/internal/launcher" @@ -47,15 +51,26 @@ type runtimeState struct { } type accountView struct { - ID string `json:"id"` - Alias string `json:"alias"` - AccountID string `json:"accountId"` - WorkspaceID string `json:"workspaceId,omitempty"` - Email string `json:"email,omitempty"` - Source string `json:"source"` - Active bool `json:"active"` - AuthenticatedAt time.Time `json:"authenticatedAt"` - LastUsedAt time.Time `json:"lastUsedAt,omitempty"` + ID string `json:"id"` + Alias string `json:"alias"` + AccountID string `json:"accountId"` + WorkspaceID string `json:"workspaceId,omitempty"` + Email string `json:"email,omitempty"` + Source string `json:"source"` + Active bool `json:"active"` + AuthenticatedAt time.Time `json:"authenticatedAt"` + LastUsedAt time.Time `json:"lastUsedAt,omitempty"` + Usage *usageView `json:"usage,omitempty"` +} + +type usageView struct { + Status string `json:"status"` + FetchedAt time.Time `json:"fetchedAt,omitempty"` + PlanType string `json:"planType,omitempty"` + RateLimits *codexusage.RateLimits `json:"rateLimits,omitempty"` + TokenUsage *codexusage.TokenUsage `json:"tokenUsage,omitempty"` + Partial []string `json:"partial,omitempty"` + Error string `json:"error,omitempty"` } func NewCommand(version string) *cobra.Command { @@ -132,6 +147,7 @@ func newAccountCommand(options *Options) *cobra.Command { newAccountImportCommand(options), newAccountAddCommand(options, false), newAccountListCommand(options), + newAccountUsageCommand(options), newAccountShowCommand(options), newAccountRenameCommand(options), newAccountRemoveCommand(options), @@ -208,10 +224,15 @@ func newAccountAddCommand(options *Options, reauth bool) *cobra.Command { } func newAccountListCommand(options *Options) *cobra.Command { - return &cobra.Command{ + var refresh bool + var cached bool + command := &cobra.Command{ Use: "list", - Short: "List saved accounts", + Short: "List saved accounts and their usage", RunE: func(*cobra.Command, []string) error { + if refresh && cached { + return errors.New("--refresh and --cached cannot be used together") + } runtime, err := options.loadRuntime(false) if err != nil { return err @@ -221,9 +242,41 @@ func newAccountListCommand(options *Options) *cobra.Command { return err } state, _ := appstate.Load(runtime.paths.State) + cache, err := runtime.usageService(options.Version).Cached() + if err != nil { + return err + } + refreshIDs := make([]string, 0, len(data.Profiles)) + if !cached { + for _, profile := range data.Profiles { + snapshot, ok := cache.Profiles[profile.ID] + if refresh || !ok || usageStatus(snapshot.FetchedAt, time.Now()) != "fresh" { + refreshIDs = append(refreshIDs, profile.ID) + } + } + } + refreshErrors := map[string]string{} + if len(refreshIDs) > 0 { + results, refreshErr := runtime.usageService(options.Version).Refresh(context.Background(), refreshIDs) + if refreshErr != nil { + for _, id := range refreshIDs { + refreshErrors[id] = refreshErr.Error() + } + } else { + for id, result := range results { + refreshErrors[id] = result.Error + } + } + cache, err = runtime.usageService(options.Version).Cached() + if err != nil { + return err + } + } views := make([]accountView, 0, len(data.Profiles)) for _, profile := range data.Profiles { - views = append(views, toView(profile, profile.ID == state.ActiveProfileID)) + view := toView(profile, profile.ID == state.ActiveProfileID) + view.Usage = usageFromCache(cache.Profiles, profile.ID, refreshErrors[profile.ID], time.Now()) + views = append(views, view) } if options.JSON { return options.render(views, "") @@ -232,6 +285,8 @@ func newAccountListCommand(options *Options) *cobra.Command { _, _ = fmt.Fprintln(options.Output, "No accounts saved.") return nil } + writer := tabwriter.NewWriter(options.Output, 0, 4, 2, ' ', 0) + _, _ = fmt.Fprintln(writer, "\tALIAS\tIDENTITY\tPLAN\tLIMITS\tTOKENS\tUPDATED") for _, view := range views { marker := " " if view.Active { @@ -241,11 +296,114 @@ func newAccountListCommand(options *Options) *cobra.Command { if identity == "" { identity = view.AccountID } - _, _ = fmt.Fprintf(options.Output, "%s %-20s %s\n", marker, view.Alias, identity) + plan, limits, tokens, updated := summarizeUsage(view.Usage, time.Now()) + _, _ = fmt.Fprintf(writer, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n", marker, view.Alias, identity, plan, limits, tokens, updated) + } + return writer.Flush() + }, + } + command.Flags().BoolVar(&refresh, "refresh", false, "refresh every account before listing") + command.Flags().BoolVar(&cached, "cached", false, "show cached usage without contacting Codex services") + return command +} + +func newAccountUsageCommand(options *Options) *cobra.Command { + var all bool + var cached bool + command := &cobra.Command{ + Use: "usage [alias]", + Short: "Show account limits and token usage without switching", + Args: cobra.MaximumNArgs(1), + RunE: func(_ *cobra.Command, args []string) error { + if all && len(args) > 0 { + return errors.New("an alias and --all cannot be used together") + } + runtime, err := options.loadRuntime(false) + if err != nil { + return err + } + data, err := runtime.manager.Load() + if err != nil { + return err + } + state, _ := appstate.Load(runtime.paths.State) + profiles := make([]vault.Profile, 0, len(data.Profiles)) + switch { + case all: + profiles = append(profiles, data.Profiles...) + case len(args) == 1: + profile, findErr := data.Find(args[0]) + if findErr != nil { + return findErr + } + profiles = append(profiles, *profile) + default: + if state.ActiveProfileID == "" { + return errors.New("no managed account is active; pass an alias or --all") + } + profile, findErr := data.Find(state.ActiveProfileID) + if findErr != nil { + return errors.New("the active account is unmanaged; pass a saved alias") + } + profiles = append(profiles, *profile) + } + if len(profiles) == 0 { + return errors.New("no accounts saved") + } + refreshErrors := map[string]string{} + if !cached { + ids := make([]string, 0, len(profiles)) + for _, profile := range profiles { + ids = append(ids, profile.ID) + } + results, refreshErr := runtime.usageService(options.Version).Refresh(context.Background(), ids) + if refreshErr != nil { + if len(profiles) == 1 { + return refreshErr + } + for _, id := range ids { + refreshErrors[id] = refreshErr.Error() + } + } else { + for id, result := range results { + refreshErrors[id] = result.Error + } + } + } + cache, err := runtime.usageService(options.Version).Cached() + if err != nil { + return err + } + views := make([]accountView, 0, len(profiles)) + for _, profile := range profiles { + view := toView(profile, profile.ID == state.ActiveProfileID) + view.Usage = usageFromCache(cache.Profiles, 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) + } + return errors.New("no cached usage is available; retry without --cached") + } + views = append(views, view) + } + if options.JSON { + if len(views) == 1 { + return options.render(views[0], "") + } + return options.render(views, "") + } + for index, view := range views { + if index > 0 { + _, _ = fmt.Fprintln(options.Output) + } + _, _ = fmt.Fprintln(options.Output, formatUsage(view, time.Now())) } return nil }, } + command.Flags().BoolVar(&all, "all", false, "show usage for every saved account") + command.Flags().BoolVar(&cached, "cached", false, "show cached usage without contacting Codex services") + return command } func newAccountShowCommand(options *Options) *cobra.Command { @@ -698,6 +856,18 @@ func (runtime runtimeState) switcher() switcher.Service { return switcher.Service{Home: runtime.home, Paths: runtime.paths, Vault: runtime.manager} } +func (runtime runtimeState) usageService(version string) accountusage.Service { + return accountusage.Service{ + Home: runtime.home, + Paths: runtime.paths, + Vault: runtime.manager, + Runner: codexusage.Runner{ + Binary: runtime.bin, + ClientVersion: version, + }, + } +} + func (options *Options) currentView() (accountView, error) { runtime, err := options.loadRuntime(false) if err != nil { @@ -769,6 +939,233 @@ 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] + if !ok { + return &usageView{Status: "unavailable", Error: queryError} + } + return &usageView{ + Status: usageStatus(snapshot.FetchedAt, now), + FetchedAt: snapshot.FetchedAt, + PlanType: snapshot.PlanType, + RateLimits: snapshot.RateLimits, + TokenUsage: snapshot.TokenUsage, + Partial: snapshot.Partial, + Error: queryError, + } +} + +func usageStatus(fetchedAt, now time.Time) string { + if fetchedAt.IsZero() || now.Sub(fetchedAt) > time.Minute { + return "stale" + } + return "fresh" +} + +func summarizeUsage(usage *usageView, now time.Time) (string, string, string, string) { + if usage == nil || usage.Status == "unavailable" { + return "-", "unavailable", "-", "-" + } + plan := usage.PlanType + if plan == "" { + plan = "-" + } + limits := "-" + if main := mainRateLimit(usage.RateLimits); main != nil { + parts := make([]string, 0, 2) + if main.Primary != nil { + parts = append(parts, compactWindow(main.Primary)) + } + if main.Secondary != nil { + parts = append(parts, compactWindow(main.Secondary)) + } + if len(parts) > 0 { + limits = strings.Join(parts, " · ") + } + } + tokens := "-" + if usage.TokenUsage != nil && usage.TokenUsage.Summary.LifetimeTokens != nil { + tokens = compactNumber(*usage.TokenUsage.Summary.LifetimeTokens) + } + updated := relativeTime(usage.FetchedAt, now) + if usage.Status == "stale" { + updated += " (stale)" + } + if usage.Error != "" { + updated += " (error)" + } + return plan, limits, tokens, updated +} + +func formatUsage(view accountView, now time.Time) string { + var output strings.Builder + output.WriteString(formatView(view)) + if view.Active { + output.WriteString(" [active]") + } + usage := view.Usage + if usage == nil || usage.Status == "unavailable" { + output.WriteString("\nUsage: unavailable") + if usage != nil && usage.Error != "" { + output.WriteString("\nError: ") + output.WriteString(usage.Error) + } + return output.String() + } + plan := usage.PlanType + if plan == "" { + plan = "unknown" + } + fmt.Fprintf(&output, "\nPlan: %s", plan) + fmt.Fprintf(&output, "\nUpdated: %s (%s)", usage.FetchedAt.Local().Format(time.RFC3339), usage.Status) + if usage.RateLimits != nil { + output.WriteString("\nLimits:") + limits := usage.RateLimits.RateLimitsByLimitID + if len(limits) == 0 { + limits = map[string]codexusage.RateLimitSnapshot{"codex": usage.RateLimits.RateLimits} + } + keys := make([]string, 0, len(limits)) + for key := range limits { + keys = append(keys, key) + } + sort.Strings(keys) + for _, key := range keys { + limit := limits[key] + label := key + if limit.LimitName != nil && *limit.LimitName != "" { + label = *limit.LimitName + } + windows := make([]string, 0, 2) + if limit.Primary != nil { + windows = append(windows, detailedWindow(limit.Primary, now)) + } + if limit.Secondary != nil { + windows = append(windows, detailedWindow(limit.Secondary, now)) + } + value := "unavailable" + if len(windows) > 0 { + value = strings.Join(windows, "; ") + } + fmt.Fprintf(&output, "\n %s: %s", label, value) + } + } + if usage.TokenUsage != nil { + summary := usage.TokenUsage.Summary + parts := make([]string, 0, 5) + if summary.LifetimeTokens != nil { + parts = append(parts, compactNumber(*summary.LifetimeTokens)+" lifetime") + } + if summary.PeakDailyTokens != nil { + parts = append(parts, compactNumber(*summary.PeakDailyTokens)+" peak/day") + } + if summary.CurrentStreakDays != nil { + parts = append(parts, fmt.Sprintf("%dd current streak", *summary.CurrentStreakDays)) + } + if summary.LongestStreakDays != nil { + parts = append(parts, fmt.Sprintf("%dd longest streak", *summary.LongestStreakDays)) + } + if summary.LongestRunningTurnSec != nil { + parts = append(parts, (time.Duration(*summary.LongestRunningTurnSec)*time.Second).String()+" longest turn") + } + if len(parts) > 0 { + output.WriteString("\nTokens: ") + output.WriteString(strings.Join(parts, "; ")) + } + } + if len(usage.Partial) > 0 { + output.WriteString("\nPartial: ") + output.WriteString(strings.Join(usage.Partial, ", ")) + } + if usage.Error != "" { + output.WriteString("\nWarning: ") + output.WriteString(usage.Error) + } + return output.String() +} + +func mainRateLimit(limits *codexusage.RateLimits) *codexusage.RateLimitSnapshot { + if limits == nil { + return nil + } + if value, ok := limits.RateLimitsByLimitID["codex"]; ok { + copy := value + return © + } + copy := limits.RateLimits + return © +} + +func compactWindow(window *codexusage.RateLimitWindow) string { + duration := "limit" + if window.WindowDurationMins != nil { + duration = compactMinutes(*window.WindowDurationMins) + } + return fmt.Sprintf("%s %d%%", duration, window.UsedPercent) +} + +func detailedWindow(window *codexusage.RateLimitWindow, now time.Time) string { + value := compactWindow(window) + " used" + if window.ResetsAt != nil { + reset := time.Unix(*window.ResetsAt, 0) + if reset.After(now) { + value += ", resets in " + compactDuration(reset.Sub(now)) + } else { + value += ", reset pending" + } + } + return value +} + +func compactMinutes(minutes int64) string { + if minutes%(60*24) == 0 { + return fmt.Sprintf("%dd", minutes/(60*24)) + } + if minutes%60 == 0 { + return fmt.Sprintf("%dh", minutes/60) + } + return fmt.Sprintf("%dm", minutes) +} + +func compactDuration(duration time.Duration) string { + if duration < 0 { + duration = 0 + } + if duration >= 24*time.Hour { + return fmt.Sprintf("%dd %dh", int(duration/(24*time.Hour)), int(duration%(24*time.Hour)/time.Hour)) + } + if duration >= time.Hour { + return fmt.Sprintf("%dh %dm", int(duration/time.Hour), int(duration%time.Hour/time.Minute)) + } + return fmt.Sprintf("%dm", int(duration/time.Minute)) +} + +func compactNumber(value int64) string { + switch { + case value >= 1_000_000_000: + return strings.TrimRight(strings.TrimRight(fmt.Sprintf("%.1f", float64(value)/1_000_000_000), "0"), ".") + "B" + case value >= 1_000_000: + return strings.TrimRight(strings.TrimRight(fmt.Sprintf("%.1f", float64(value)/1_000_000), "0"), ".") + "M" + case value >= 1_000: + return strings.TrimRight(strings.TrimRight(fmt.Sprintf("%.1f", float64(value)/1_000), "0"), ".") + "K" + default: + return strconv.FormatInt(value, 10) + } +} + +func relativeTime(value, now time.Time) string { + if value.IsZero() { + return "never" + } + age := now.Sub(value) + if age < 0 { + age = 0 + } + if age < time.Minute { + return "just now" + } + return compactDuration(age) + " ago" +} + func (options *Options) render(value any, text string) error { if options.JSON { encoder := json.NewEncoder(options.Output) diff --git a/internal/app/usage_test.go b/internal/app/usage_test.go new file mode 100644 index 0000000..a652a3a --- /dev/null +++ b/internal/app/usage_test.go @@ -0,0 +1,58 @@ +package app + +import ( + "strings" + "testing" + "time" + + "github.com/SilkageNet/codex-switch/internal/codexusage" +) + +func TestSummarizeUsage(t *testing.T) { + now := time.Date(2026, 8, 21, 12, 0, 0, 0, time.UTC) + primaryDuration := int64(300) + secondaryDuration := int64(10080) + lifetime := int64(1_234_567) + view := &usageView{ + Status: "fresh", + FetchedAt: now.Add(-10 * time.Second), + PlanType: "pro", + RateLimits: &codexusage.RateLimits{RateLimitsByLimitID: map[string]codexusage.RateLimitSnapshot{ + "codex": { + Primary: &codexusage.RateLimitWindow{UsedPercent: 21, WindowDurationMins: &primaryDuration}, + Secondary: &codexusage.RateLimitWindow{UsedPercent: 81, WindowDurationMins: &secondaryDuration}, + }, + }}, + TokenUsage: &codexusage.TokenUsage{Summary: codexusage.TokenUsageSummary{LifetimeTokens: &lifetime}}, + } + plan, limits, tokens, updated := summarizeUsage(view, now) + if plan != "pro" || limits != "5h 21% · 7d 81%" || tokens != "1.2M" || updated != "just now" { + t.Fatalf("unexpected summary: %q %q %q %q", plan, limits, tokens, updated) + } +} + +func TestFormatUsageShowsDetailedWindows(t *testing.T) { + now := time.Date(2026, 8, 21, 12, 0, 0, 0, time.UTC) + duration := int64(300) + reset := now.Add(2*time.Hour + 15*time.Minute).Unix() + name := "Codex" + view := accountView{ + Alias: "work", + Email: "person@example.com", + Active: true, + Usage: &usageView{ + Status: "fresh", + FetchedAt: now, + PlanType: "plus", + RateLimits: &codexusage.RateLimits{RateLimitsByLimitID: map[string]codexusage.RateLimitSnapshot{ + "codex": {LimitName: &name, Primary: &codexusage.RateLimitWindow{UsedPercent: 33, WindowDurationMins: &duration, ResetsAt: &reset}}, + }}, + }, + } + output := formatUsage(view, now) + for _, expected := range []string{"work (person@example.com) [active]", "Plan: plus", "Codex: 5h 33% used, resets in 2h 15m"} { + if !strings.Contains(output, expected) { + t.Fatalf("output %q does not contain %q", output, expected) + } + } +} diff --git a/internal/codexusage/runner.go b/internal/codexusage/runner.go new file mode 100644 index 0000000..d0e0f6a --- /dev/null +++ b/internal/codexusage/runner.go @@ -0,0 +1,287 @@ +package codexusage + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/SilkageNet/codex-switch/internal/atomicfile" + "github.com/SilkageNet/codex-switch/internal/authschema" +) + +const maxProtocolMessage = 8 << 20 + +type Runner struct { + Binary string + ClientVersion string + command func(context.Context, string, ...string) *exec.Cmd +} + +type rpcRequest struct { + Method string `json:"method"` + ID *int `json:"id,omitempty"` + Params any `json:"params,omitempty"` +} + +type rpcResponse struct { + ID json.RawMessage `json:"id"` + Result json.RawMessage `json:"result"` + Error *rpcError `json:"error"` +} + +type rpcError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +func (runner Runner) Query(ctx context.Context, auth json.RawMessage) (Snapshot, json.RawMessage, error) { + if runner.Binary == "" { + return Snapshot{}, nil, errors.New("official Codex executable not found; install Codex or pass --codex-binary") + } + document, err := authschema.Parse(auth) + if err != nil { + return Snapshot{}, nil, fmt.Errorf("validate saved account: %w", err) + } + temporaryHome, err := os.MkdirTemp("", "codex-switch-usage-*") + if err != nil { + return Snapshot{}, nil, fmt.Errorf("create temporary CODEX_HOME: %w", err) + } + defer func() { _ = os.RemoveAll(temporaryHome) }() + if err := os.Chmod(temporaryHome, 0o700); err != nil { + return Snapshot{}, nil, fmt.Errorf("protect temporary CODEX_HOME: %w", err) + } + if err := atomicfile.Write(filepath.Join(temporaryHome, "config.toml"), []byte("cli_auth_credentials_store = \"file\"\n"), 0o600); err != nil { + return Snapshot{}, nil, err + } + if err := atomicfile.Write(filepath.Join(temporaryHome, "auth.json"), append(append([]byte(nil), document.Raw...), '\n'), 0o600); err != nil { + return Snapshot{}, nil, err + } + + commandFactory := runner.command + if commandFactory == nil { + commandFactory = exec.CommandContext + } + command := commandFactory(ctx, runner.Binary, "app-server") + environment := command.Env + if environment == nil { + environment = os.Environ() + } + command.Env = withEnvironment(environment, "CODEX_HOME", temporaryHome) + stdin, err := command.StdinPipe() + if err != nil { + return Snapshot{}, nil, err + } + stdout, err := command.StdoutPipe() + if err != nil { + return Snapshot{}, nil, err + } + var stderr limitedBuffer + command.Stderr = &stderr + if err := command.Start(); err != nil { + return Snapshot{}, nil, fmt.Errorf("start Codex app server: %w", err) + } + waited := false + defer func() { + _ = stdin.Close() + if !waited && command.Process != nil { + _ = command.Process.Kill() + _ = command.Wait() + } + }() + + encoder := json.NewEncoder(stdin) + scanner := bufio.NewScanner(stdout) + scanner.Buffer(make([]byte, 64<<10), maxProtocolMessage) + version := runner.ClientVersion + if version == "" { + version = "dev" + } + if err := sendRequest(encoder, 0, "initialize", map[string]any{ + "clientInfo": map[string]string{"name": "codex_switch", "title": "codex-switch", "version": version}, + }); err != nil { + return Snapshot{}, nil, err + } + initialize, err := readResponse(scanner, 0) + if err != nil { + return Snapshot{}, nil, protocolError("initialize", err, stderr.String()) + } + if initialize.Error != nil { + return Snapshot{}, nil, responseError("initialize", initialize.Error) + } + 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 + } + } + responses, err := readResponses(scanner, 1, 2, 3) + if err != nil { + return Snapshot{}, nil, protocolError("read account usage", err, stderr.String()) + } + _ = stdin.Close() + if err := command.Wait(); err != nil { + waited = true + if ctx.Err() != nil { + return Snapshot{}, nil, fmt.Errorf("query account usage: %w", ctx.Err()) + } + 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) + } + 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") + } + snapshot := Snapshot{FetchedAt: time.Now().UTC(), PlanType: account.Account.PlanType} + if response := responses[2]; response.Error != nil { + snapshot.Partial = append(snapshot.Partial, "rateLimits") + } else { + var value RateLimits + if err := json.Unmarshal(response.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 { + snapshot.Partial = append(snapshot.Partial, "tokenUsage") + } else { + var value TokenUsage + if err := json.Unmarshal(response.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, nil +} + +func sendRequest(encoder *json.Encoder, id int, method string, params any) error { + if err := encoder.Encode(rpcRequest{Method: method, ID: &id, Params: params}); err != nil { + return fmt.Errorf("send %s request: %w", method, err) + } + return nil +} + +func readResponse(scanner *bufio.Scanner, expected int) (rpcResponse, error) { + responses, err := readResponses(scanner, expected) + return responses[expected], err +} + +func readResponses(scanner *bufio.Scanner, expected ...int) (map[int]rpcResponse, error) { + wanted := make(map[int]bool, len(expected)) + for _, id := range expected { + wanted[id] = true + } + responses := make(map[int]rpcResponse, len(expected)) + for len(responses) < len(wanted) && scanner.Scan() { + line := bytes.TrimSpace(scanner.Bytes()) + if len(line) == 0 { + continue + } + var response rpcResponse + if err := json.Unmarshal(line, &response); err != nil { + return nil, fmt.Errorf("decode JSON-RPC message: %w", err) + } + if len(response.ID) == 0 || bytes.Equal(response.ID, []byte("null")) { + continue + } + id, err := strconv.Atoi(strings.Trim(string(response.ID), `"`)) + if err != nil || !wanted[id] { + continue + } + responses[id] = response + } + if err := scanner.Err(); err != nil { + return nil, err + } + if len(responses) != len(wanted) { + return nil, io.ErrUnexpectedEOF + } + return responses, nil +} + +func responseError(method string, rpcErr *rpcError) error { + return fmt.Errorf("%s failed (%d): %s", method, rpcErr.Code, rpcErr.Message) +} + +func protocolError(operation string, err error, stderr string) error { + stderr = strings.TrimSpace(stderr) + if stderr == "" { + return fmt.Errorf("%s: %w", operation, err) + } + return fmt.Errorf("%s: %w (Codex app-server diagnostics redacted)", operation, err) +} + +func withEnvironment(environment []string, key, value string) []string { + prefix := key + "=" + result := make([]string, 0, len(environment)+1) + for _, entry := range environment { + if !strings.HasPrefix(entry, prefix) { + result = append(result, entry) + } + } + return append(result, prefix+value) +} + +type limitedBuffer struct { + buffer bytes.Buffer +} + +func (buffer *limitedBuffer) Write(data []byte) (int, error) { + const limit = 16 << 10 + original := len(data) + remaining := limit - buffer.buffer.Len() + if remaining > 0 { + if len(data) > remaining { + data = data[:remaining] + } + _, _ = buffer.buffer.Write(data) + } + return original, nil +} + +func (buffer *limitedBuffer) String() string { + return buffer.buffer.String() +} diff --git a/internal/codexusage/runner_test.go b/internal/codexusage/runner_test.go new file mode 100644 index 0000000..25914fb --- /dev/null +++ b/internal/codexusage/runner_test.go @@ -0,0 +1,124 @@ +package codexusage + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" +) + +func TestRunnerQueriesUsageAndReturnsRefreshedCredentials(t *testing.T) { + runner := helperRunner(t, false) + snapshot, updated, err := runner.Query(context.Background(), testAuth("account-a", "refresh-old", "2026-08-20T00:00:00Z")) + if err != nil { + t.Fatal(err) + } + if snapshot.PlanType != "pro" || snapshot.RateLimits == nil || snapshot.TokenUsage == nil { + t.Fatalf("unexpected snapshot: %#v", snapshot) + } + if snapshot.MainRateLimit() == nil || snapshot.MainRateLimit().Primary.UsedPercent != 21 { + t.Fatalf("unexpected main rate limit: %#v", snapshot.MainRateLimit()) + } + if snapshot.TokenUsage.Summary.LifetimeTokens == nil || *snapshot.TokenUsage.Summary.LifetimeTokens != 1234567 { + t.Fatalf("unexpected token usage: %#v", snapshot.TokenUsage) + } + var wire struct { + Tokens struct { + Refresh string `json:"refresh_token"` + } `json:"tokens"` + } + if err := json.Unmarshal(updated, &wire); err != nil { + t.Fatal(err) + } + if wire.Tokens.Refresh != "refresh-new" { + t.Fatalf("refreshed credentials were not returned: %s", wire.Tokens.Refresh) + } +} + +func TestRunnerAllowsOneUnsupportedUsageMethod(t *testing.T) { + runner := helperRunner(t, true) + snapshot, _, err := runner.Query(context.Background(), testAuth("account-a", "refresh-old", "2026-08-20T00:00:00Z")) + if err != nil { + t.Fatal(err) + } + if snapshot.RateLimits == nil || snapshot.TokenUsage != nil || len(snapshot.Partial) != 1 || snapshot.Partial[0] != "tokenUsage" { + t.Fatalf("unexpected partial snapshot: %#v", snapshot) + } +} + +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") { + t.Fatalf("unexpected protocol error: %v", err) + } +} + +func helperRunner(t *testing.T, partial bool) Runner { + t.Helper() + return Runner{ + Binary: os.Args[0], + ClientVersion: "test", + command: func(ctx context.Context, binary string, args ...string) *exec.Cmd { + command := exec.CommandContext(ctx, binary, "-test.run=TestCodexUsageHelperProcess", "--") + command.Env = append(os.Environ(), "CODEX_USAGE_HELPER=1") + if partial { + command.Env = append(command.Env, "CODEX_USAGE_PARTIAL=1") + } + return command + }, + } +} + +func TestCodexUsageHelperProcess(t *testing.T) { + if os.Getenv("CODEX_USAGE_HELPER") != "1" { + return + } + encoder := json.NewEncoder(os.Stdout) + scanner := bufio.NewScanner(os.Stdin) + for scanner.Scan() { + var request struct { + Method string `json:"method"` + ID json.RawMessage `json:"id"` + } + if json.Unmarshal(scanner.Bytes(), &request) != nil || len(request.ID) == 0 { + continue + } + id := string(request.ID) + switch request.Method { + case "initialize": + _ = encoder.Encode(map[string]any{"id": json.RawMessage(id), "result": map[string]any{}}) + case "account/read": + _ = encoder.Encode(map[string]any{"id": json.RawMessage(id), "result": map[string]any{ + "account": map[string]any{"type": "chatgpt", "email": "a@example.com", "planType": "pro"}, + "requiresOpenaiAuth": true, + }}) + case "account/rateLimits/read": + _ = 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}}}, + }}) + case "account/usage/read": + if os.Getenv("CODEX_USAGE_PARTIAL") == "1" { + _ = encoder.Encode(map[string]any{"id": json.RawMessage(id), "error": map[string]any{"code": -32601, "message": "method not found"}}) + } else { + lifetime := int64(1234567) + _ = encoder.Encode(map[string]any{"id": json.RawMessage(id), "result": map[string]any{ + "summary": map[string]any{"lifetimeTokens": lifetime}, + "dailyUsageBuckets": []map[string]any{{"startDate": "2026-08-20", "tokens": 99}}, + }}) + } + home := os.Getenv("CODEX_HOME") + _ = os.WriteFile(filepath.Join(home, "auth.json"), testAuth("account-a", "refresh-new", "2026-08-20T01:00:00Z"), 0o600) + } + } + os.Exit(0) +} + +func testAuth(account, refresh, lastRefresh string) []byte { + return []byte(fmt.Sprintf(`{"auth_mode":"chatgpt","tokens":{"id_token":%q,"access_token":%q,"refresh_token":%q,"account_id":%q},"last_refresh":%q}`, "id-"+refresh, "access-"+refresh, refresh, account, lastRefresh)) +} diff --git a/internal/codexusage/types.go b/internal/codexusage/types.go new file mode 100644 index 0000000..9e052b7 --- /dev/null +++ b/internal/codexusage/types.go @@ -0,0 +1,104 @@ +package codexusage + +import "time" + +type Snapshot struct { + FetchedAt time.Time `json:"fetchedAt"` + PlanType string `json:"planType,omitempty"` + RateLimits *RateLimits `json:"rateLimits,omitempty"` + TokenUsage *TokenUsage `json:"tokenUsage,omitempty"` + Partial []string `json:"partial,omitempty"` +} + +type AccountResponse struct { + Account *Account `json:"account"` + RequiresOpenAIAuth bool `json:"requiresOpenaiAuth"` +} + +type Account struct { + Type string `json:"type"` + Email *string `json:"email,omitempty"` + PlanType string `json:"planType,omitempty"` +} + +type RateLimits struct { + RateLimits RateLimitSnapshot `json:"rateLimits"` + RateLimitsByLimitID map[string]RateLimitSnapshot `json:"rateLimitsByLimitId,omitempty"` + RateLimitResetCredits *RateLimitResetCreditsSummary `json:"rateLimitResetCredits,omitempty"` +} + +type RateLimitSnapshot struct { + LimitID *string `json:"limitId,omitempty"` + LimitName *string `json:"limitName,omitempty"` + PlanType *string `json:"planType,omitempty"` + Primary *RateLimitWindow `json:"primary,omitempty"` + Secondary *RateLimitWindow `json:"secondary,omitempty"` + Credits *CreditsSnapshot `json:"credits,omitempty"` + IndividualLimit *SpendControlLimitSnapshot `json:"individualLimit,omitempty"` + SpendControlReached *bool `json:"spendControlReached,omitempty"` + RateLimitReachedType *string `json:"rateLimitReachedType,omitempty"` +} + +type RateLimitWindow struct { + UsedPercent int `json:"usedPercent"` + WindowDurationMins *int64 `json:"windowDurationMins,omitempty"` + ResetsAt *int64 `json:"resetsAt,omitempty"` +} + +type CreditsSnapshot struct { + HasCredits bool `json:"hasCredits"` + Unlimited bool `json:"unlimited"` + Balance *string `json:"balance,omitempty"` +} + +type SpendControlLimitSnapshot struct { + Limit string `json:"limit"` + Used string `json:"used"` + RemainingPercent int `json:"remainingPercent"` + ResetsAt int64 `json:"resetsAt"` +} + +type RateLimitResetCreditsSummary struct { + AvailableCount int64 `json:"availableCount"` + Credits []RateLimitResetCredit `json:"credits,omitempty"` +} + +type RateLimitResetCredit struct { + ID string `json:"id"` + ResetType string `json:"resetType"` + Status string `json:"status"` + GrantedAt int64 `json:"grantedAt"` + ExpiresAt *int64 `json:"expiresAt,omitempty"` + Title *string `json:"title,omitempty"` + Description *string `json:"description,omitempty"` +} + +type TokenUsage struct { + Summary TokenUsageSummary `json:"summary"` + DailyUsageBuckets []DailyUsageBucket `json:"dailyUsageBuckets,omitempty"` +} + +type TokenUsageSummary struct { + LifetimeTokens *int64 `json:"lifetimeTokens,omitempty"` + PeakDailyTokens *int64 `json:"peakDailyTokens,omitempty"` + LongestRunningTurnSec *int64 `json:"longestRunningTurnSec,omitempty"` + CurrentStreakDays *int64 `json:"currentStreakDays,omitempty"` + LongestStreakDays *int64 `json:"longestStreakDays,omitempty"` +} + +type DailyUsageBucket struct { + StartDate string `json:"startDate"` + Tokens int64 `json:"tokens"` +} + +func (snapshot Snapshot) MainRateLimit() *RateLimitSnapshot { + if snapshot.RateLimits == nil { + return nil + } + if limit, ok := snapshot.RateLimits.RateLimitsByLimitID["codex"]; ok { + copy := limit + return © + } + copy := snapshot.RateLimits.RateLimits + return © +} diff --git a/internal/config/paths.go b/internal/config/paths.go index 2f625e0..35fbd42 100644 --- a/internal/config/paths.go +++ b/internal/config/paths.go @@ -7,10 +7,11 @@ import ( ) type Paths struct { - Root string - Vault string - State string - Journal string + Root string + Vault string + State string + Journal string + UsageCache string } func ResolvePaths(override string) (Paths, error) { @@ -30,10 +31,11 @@ func ResolvePaths(override string) (Paths, error) { return Paths{}, fmt.Errorf("resolve codex-switch home: %w", err) } return Paths{ - Root: abs, - Vault: filepath.Join(abs, "vault.v1.enc"), - State: filepath.Join(abs, "state.json"), - Journal: filepath.Join(abs, "switch.journal.json"), + Root: abs, + Vault: filepath.Join(abs, "vault.v1.enc"), + State: filepath.Join(abs, "state.json"), + Journal: filepath.Join(abs, "switch.journal.json"), + UsageCache: filepath.Join(abs, "usage-cache.v1.json"), }, nil } diff --git a/internal/usagecache/cache.go b/internal/usagecache/cache.go new file mode 100644 index 0000000..c00519b --- /dev/null +++ b/internal/usagecache/cache.go @@ -0,0 +1,49 @@ +package usagecache + +import ( + "encoding/json" + "errors" + "fmt" + "os" + + "github.com/SilkageNet/codex-switch/internal/atomicfile" + "github.com/SilkageNet/codex-switch/internal/codexusage" +) + +type Cache struct { + Version int `json:"version"` + Profiles map[string]codexusage.Snapshot `json:"profiles"` +} + +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 + } + if err != nil { + return Cache{}, err + } + var cache Cache + if err := json.Unmarshal(data, &cache); err != nil { + return Cache{}, fmt.Errorf("decode usage cache: %w", err) + } + if cache.Version != 1 { + return Cache{}, fmt.Errorf("unsupported usage cache version %d", cache.Version) + } + if cache.Profiles == nil { + cache.Profiles = map[string]codexusage.Snapshot{} + } + return cache, nil +} + +func Save(path string, cache Cache) error { + cache.Version = 1 + if cache.Profiles == nil { + cache.Profiles = map[string]codexusage.Snapshot{} + } + data, err := json.MarshalIndent(cache, "", " ") + if err != nil { + return err + } + return atomicfile.Write(path, append(data, '\n'), 0o600) +} diff --git a/internal/usagecache/cache_test.go b/internal/usagecache/cache_test.go new file mode 100644 index 0000000..3cbeb9c --- /dev/null +++ b/internal/usagecache/cache_test.go @@ -0,0 +1,43 @@ +package usagecache + +import ( + "os" + "path/filepath" + "runtime" + "testing" + "time" + + "github.com/SilkageNet/codex-switch/internal/codexusage" +) + +func TestLoadMissingAndRoundTrip(t *testing.T) { + path := filepath.Join(t.TempDir(), "nested", "usage.json") + cache, err := Load(path) + if err != nil { + t.Fatal(err) + } + if cache.Version != 1 || cache.Profiles == nil { + t.Fatalf("unexpected empty cache: %#v", cache) + } + fetched := time.Date(2026, 8, 20, 1, 2, 3, 0, time.UTC) + cache.Profiles["profile-a"] = codexusage.Snapshot{FetchedAt: fetched, PlanType: "pro"} + if err := Save(path, cache); err != nil { + t.Fatal(err) + } + loaded, err := Load(path) + if err != nil { + t.Fatal(err) + } + if loaded.Profiles["profile-a"].PlanType != "pro" || !loaded.Profiles["profile-a"].FetchedAt.Equal(fetched) { + t.Fatalf("unexpected cache round trip: %#v", loaded) + } + if runtime.GOOS != "windows" { + info, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + if info.Mode().Perm() != 0o600 { + t.Fatalf("cache mode = %o", info.Mode().Perm()) + } + } +}