Merge pull request #17 from drogers0/fix-codex-refresh-expiry

fix(codex): gate refresh on access-token expiry, not the id_token

Author: drogers0

CLAUDE.md

@@ -35,9 +35,9 @@ JSON is the default; exit codes: 0 success / 1 any-provider-failed / 2 usage err
 - [internal/providers/usagecache/](internal/providers/usagecache/) — provider-neutral 90 s file-backed usage cache. `New(provider, nowFn, warnFn)` validates the provider char-set and writes `$CACHE/aistat/usage/<provider>-v1.json` with `<provider>.cache.lock`. Warn strings include the provider name.
 - [internal/providers/multiaccount/](internal/providers/multiaccount/) — provider-neutral helpers consumed by both Claude and Codex multi-account `Fetch`: `SortAccountResults`, `RecordFetchOutcome`, `RecomputeResetAfter`, `Budget(base, perAccount, count)`.
 - [internal/providers/claude/](internal/providers/claude/) — Claude's multi-account machinery: [claude.go](internal/providers/claude/claude.go) (Fetch + FetchForSwitch + cache wiring), [profile.go](internal/providers/claude/profile.go) (identity endpoint), [refresh.go](internal/providers/claude/refresh.go) (OAuth refresh), [reconcile.go](internal/providers/claude/reconcile.go) (pure decision tree: byte-match → profile lookup → fallback), [account.go](internal/providers/claude/account.go) (`StoredAccessToken` / `StoredRefreshToken` / `StoredExpiresAt` token-parsing helpers operating on the opaque `accounts.Account`).
-- [internal/providers/codex/](internal/providers/codex/) — Codex's multi-account machinery, structurally a mirror of Claude's: [codex.go](internal/providers/codex/codex.go) (Fetch + FetchForSwitch + cache wiring + `rotateRawBlob`), [refresh.go](internal/providers/codex/refresh.go) (OAuth refresh against `https://auth.openai.com/oauth/token` with client_id confirmed via Codex binary inspection), [reconcile.go](internal/providers/codex/reconcile.go) (pure decision tree: byte-match → JWT `sub` lookup → live-unstored), [account.go](internal/providers/codex/account.go) (Codex-shaped `Stored*` helpers). Identity is the `sub` claim of the OIDC `id_token` — no network endpoint, JWT-payload decode only. Slot-vs-duration window labelling is by `limit_window_seconds`, NOT by slot position, so free-account weekly windows that land in the primary slot are not mislabelled as `five_hour`.
+- [internal/providers/codex/](internal/providers/codex/) — Codex's multi-account machinery, structurally a mirror of Claude's: [codex.go](internal/providers/codex/codex.go) (Fetch + FetchForSwitch + cache wiring + `rotateRawBlob`), [refresh.go](internal/providers/codex/refresh.go) (OAuth refresh against `https://auth.openai.com/oauth/token` with client_id confirmed via Codex binary inspection), [reconcile.go](internal/providers/codex/reconcile.go) (pure decision tree: byte-match → JWT `sub` lookup → live-unstored), [account.go](internal/providers/codex/account.go) (Codex-shaped `Stored*` helpers). `StoredExpiresAt` decodes expiry from the **access_token** JWT's `exp` (via `cred.ParseJWTExp`) — the long-lived API credential — NOT the short-lived OIDC `id_token`, which is identity-only (gating on the id_token's ~1 h `exp` made the refresh gate fire on every run). Identity is the `sub` claim of the OIDC `id_token` — no network endpoint, JWT-payload decode only. Slot-vs-duration window labelling is by `limit_window_seconds`, NOT by slot position, so free-account weekly windows that land in the primary slot are not mislabelled as `five_hour`.
 - [internal/render/](internal/render/) — `json` and `text` renderers. The JSON shape is the public contract; the text renderer is a thin presentation layer over the same model. Provider-agnostic: when `len(result.Accounts) > 0` for ANY provider the renderer emits the nested per-account view, otherwise the legacy flat form (still in use for Copilot, and as a Claude/Codex fallback).
-- [internal/cred/](internal/cred/) — credential read/write and JWT decoding. `Credential.Raw []byte` preserves the verbatim provider blob (every byte the upstream CLI wrote) so a switch re-publishes byte-for-byte. `ReadClaudeCredential` / `WriteClaudeLiveBlob` (macOS Keychain item `Claude Code-credentials` via a `runSecurity` test seam; Linux `~/.claude/.credentials.json`). `ReadCodexCredential` / `WriteCodexLiveBlob` (file-only on both OSes: `~/.codex/auth.json`, mode 0600, atomic rename + fsync). `ParseCodexIDToken` (exported here to avoid a `cred` ↔ `providers/codex` import cycle) decodes the OIDC `id_token` payload for `sub` / `email` / `exp` without signature verification.
+- [internal/cred/](internal/cred/) — credential read/write and JWT decoding. `Credential.Raw []byte` preserves the verbatim provider blob (every byte the upstream CLI wrote) so a switch re-publishes byte-for-byte. `ReadClaudeCredential` / `WriteClaudeLiveBlob` (macOS Keychain item `Claude Code-credentials` via a `runSecurity` test seam; Linux `~/.claude/.credentials.json`). `ReadCodexCredential` / `WriteCodexLiveBlob` (file-only on both OSes: `~/.codex/auth.json`, mode 0600, atomic rename + fsync). `ParseCodexIDToken` (exported here to avoid a `cred` ↔ `providers/codex` import cycle) decodes the OIDC `id_token` payload for `sub` / `email` / `exp` without signature verification; `ParseJWTExp` similarly decodes just the `exp` claim of any JWT (no signature check) and is used by `codex.StoredExpiresAt` to read the access-token expiry.
 - [internal/httpx/](internal/httpx/) — shared HTTP transport: `Doer.GetJSON` (Authorization-reserved) + `Doer.PostForm` (no Authorization by default — used by both refresh clients) sharing an unexported `setCommonHeaders` / `do` split. `do` runs a bounded retry loop (max 3 attempts) that honors `Retry-After` (capped at 10 s) on transient classifications, falling back to exponential backoff with ±20 % jitter. `Classifier` takes `*http.Response` so callers can inspect headers.
 - [internal/orchestrate/](internal/orchestrate/) — parallel fan-out across providers; one failing provider does not block the others. Preserves per-account rows on provider-level error (D8 contract).
 - [internal/testutil/](internal/testutil/) — shared test helpers.
@@ -76,7 +76,7 @@ These are the principles every change should respect. When in doubt, optimize fo
 
 These are upstream OAuth-provider behaviors aistat cannot work around without re-authentication. Both fail closed with actionable errors. **`codex login --device-auth` is the right flow on remote / headless machines** (no browser required on the host), but per upstream [code](https://github.com/openai/codex/blob/main/codex-rs/login/src/device_code_auth.rs) and [docs](https://developers.openai.com/codex/auth) device-auth uses the same persistence path and the same single-cached-login model as the browser flow — the revocation semantics below apply equally.
 
-- **Refresh-token rotation race.** OpenAI's `/oauth/token` endpoint single-uses each refresh_token; the Codex CLI rotates it on every refresh and writes the new value to `~/.codex/auth.json`. If the Codex CLI runs between aistat reading the file and aistat sending its refresh request, aistat's in-memory copy is stale and the server returns a 401 whose body reads `Your refresh token has already been used to generate a new access token.` aistat tightens this to `stale refresh token (codex CLI rotated it); retry or run codex login to recover` (matched on `already been used` in `refreshErrorMessage`). Recovery: `codex login` (or just wait for the next aistat run — the cache will catch up next pass).
+- **Refresh-token rotation race.** OpenAI's `/oauth/token` endpoint single-uses each refresh_token; the Codex CLI rotates it on every refresh and writes the new value to `~/.codex/auth.json`. If the Codex CLI runs between aistat reading the file and aistat sending its refresh request, aistat's in-memory copy is stale and the server returns a 401 whose body reads `Your refresh token has already been used to generate a new access token.` aistat tightens this to `stale refresh token (codex CLI rotated it); retry or run codex login to recover` (matched on `already been used` in `refreshErrorMessage`). Recovery: `codex login` (or just wait for the next aistat run — the cache will catch up next pass). Note: since `StoredExpiresAt` now gates on the long-lived access_token's `exp`, routine reporting no longer triggers a refresh on every run — only when the access token is genuinely near expiry — so this race is hit far less often.
 - **Switch-side token revocation.** When a new account logs in on the same OAuth client via the browser flow (`codex logout && codex login`), OpenAI's server invalidates the previous account's tokens. aistat's `switch` re-publishes the stored blob byte-for-byte, but server-revoked tokens stay revoked — the next usage call returns a 401 whose body carries either `token_revoked` or `token_invalidated` (OpenAI uses both interchangeably for this condition). aistat tightens both to `tokens revoked by upstream (likely a codex login for another account); run codex login to recover` (matched by `isRevokedTokenErr`). Recovery: `codex login` for the now-active account.
 
 ## Working in this repo

internal/cred/claude.go

@@ -27,8 +27,12 @@ var ErrClaudeWriteUnsupported = errors.New("writing Claude live credential is no
 type Credential struct {
 	AccessToken  string
 	RefreshToken string
-	ExpiresAt    int64 // ms since epoch; 0 if absent
-	Raw          []byte
+	// ExpiresAt is ms since epoch; 0 if absent. Claude only — populated from the
+	// claudeAiOauth.expiresAt field. Codex leaves this 0 (its auth.json has no
+	// expiry field); the codex refresh gate decodes the access-token JWT exp on
+	// demand in codex.StoredExpiresAt.
+	ExpiresAt int64
+	Raw       []byte
 }
 
 // parseClaudeCredFull parses the JSON payload used by both the macOS Keychain

internal/cred/codex.go

@@ -64,19 +64,47 @@ func ParseCodexIDToken(idToken string) (sub, email string, expSec int64, err err
 	return claims.Sub, claims.Email, int64(claims.Exp), nil
 }
 
+// ParseJWTExp decodes the exp claim (seconds since epoch) from a JWT payload
+// without verifying the signature — same posture as ParseCodexIDToken (the Codex
+// CLI already accepted the token). Returns (0, false) when token is not a
+// parseable 3-segment JWT or the exp claim is absent/zero. Used by
+// codex.StoredExpiresAt to read the access-token expiry that drives the refresh
+// gate.
+func ParseJWTExp(token string) (expSec int64, ok bool) {
+	parts := strings.Split(token, ".")
+	if len(parts) != 3 || parts[0] == "" || parts[1] == "" || parts[2] == "" {
+		return 0, false
+	}
+	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
+	if err != nil {
+		return 0, false
+	}
+	var claims struct {
+		Exp float64 `json:"exp"` // JSON number; may be integer or float
+	}
+	if err := json.Unmarshal(payload, &claims); err != nil {
+		return 0, false
+	}
+	if claims.Exp <= 0 {
+		return 0, false
+	}
+	return int64(claims.Exp), true
+}
+
 // rawCodexAuth is the minimal shape of ~/.codex/auth.json for credential extraction.
 type rawCodexAuth struct {
 	Tokens struct {
 		AccessToken  string `json:"access_token"`
 		RefreshToken string `json:"refresh_token"`
-		IDToken      string `json:"id_token"`
 	} `json:"tokens"`
 }
 
 // parseCodexCredFull parses the JSON payload of ~/.codex/auth.json.
 // access_token is required; its absence returns ErrCodexTokenNotFound.
-// id_token is optional: if absent or malformed, ExpiresAt is 0.
-// Raw is set to bytes.Clone(data) so the caller's buffer can be reused.
+// Credential.ExpiresAt is always 0 for codex (auth.json has no expiry field);
+// the refresh gate decodes the access-token JWT exp on demand in
+// codex.StoredExpiresAt. Raw is set to bytes.Clone(data) so the caller's buffer
+// can be reused.
 func parseCodexCredFull(data []byte) (Credential, error) {
 	var raw rawCodexAuth
 	if err := json.Unmarshal(data, &raw); err != nil {
@@ -86,20 +114,9 @@ func parseCodexCredFull(data []byte) (Credential, error) {
 		return Credential{}, ErrCodexTokenNotFound
 	}
 
-	var expiresAt int64
-	if raw.Tokens.IDToken != "" {
-		if _, _, expSec, err := ParseCodexIDToken(raw.Tokens.IDToken); err == nil {
-			expiresAt = expSec * 1000
-		}
-		// malformed id_token → ExpiresAt stays 0; not an error for ReadCodexCredential.
-		// T3 reconcile calls ParseCodexIDToken independently on the live credential
-		// to extract sub/email for identity; it will surface the error there if needed.
-	}
-
 	return Credential{
 		AccessToken:  raw.Tokens.AccessToken,
 		RefreshToken: raw.Tokens.RefreshToken,
-		ExpiresAt:    expiresAt,
 		Raw:          bytes.Clone(data),
 	}, nil
 }

internal/cred/codex_test.go

@@ -193,6 +193,62 @@ func TestParseCodexIDToken(t *testing.T) {
 	}
 }
 
+// --- ParseJWTExp tests ---
+
+func TestParseJWTExp(t *testing.T) {
+	tests := []struct {
+		name    string
+		run     func(t *testing.T) (int64, bool)
+		wantExp int64
+		wantOK  bool
+	}{
+		{"present", func(t *testing.T) (int64, bool) {
+			return ParseJWTExp(makeTestJWT(t, `{"sub":"u1","exp":1700000000}`))
+		}, 1700000000, true},
+		{"past exp still returned", func(t *testing.T) (int64, bool) {
+			return ParseJWTExp(makeTestJWT(t, `{"exp":1000000000}`))
+		}, 1000000000, true},
+		{"float exp truncated", func(t *testing.T) (int64, bool) {
+			return ParseJWTExp(makeTestJWT(t, `{"exp":1700000000.9}`))
+		}, 1700000000, true},
+		{"exp absent", func(t *testing.T) (int64, bool) {
+			return ParseJWTExp(makeTestJWT(t, `{"sub":"u1"}`))
+		}, 0, false},
+		{"exp zero", func(t *testing.T) (int64, bool) {
+			return ParseJWTExp(makeTestJWT(t, `{"exp":0}`))
+		}, 0, false},
+		{"exp negative", func(t *testing.T) (int64, bool) {
+			return ParseJWTExp(makeTestJWT(t, `{"exp":-1}`))
+		}, 0, false},
+		{"opaque non-jwt", func(t *testing.T) (int64, bool) {
+			return ParseJWTExp("test-access-token-abc")
+		}, 0, false},
+		{"two segments", func(t *testing.T) (int64, bool) {
+			return ParseJWTExp("header.payload")
+		}, 0, false},
+		{"empty segment", func(t *testing.T) (int64, bool) {
+			return ParseJWTExp("a..c")
+		}, 0, false},
+		{"bad base64 payload", func(t *testing.T) (int64, bool) {
+			return ParseJWTExp("header.!!!.sig")
+		}, 0, false},
+		{"bad json payload", func(t *testing.T) (int64, bool) {
+			return ParseJWTExp("header." + base64.RawURLEncoding.EncodeToString([]byte("not json")) + ".sig")
+		}, 0, false},
+		{"empty", func(t *testing.T) (int64, bool) {
+			return ParseJWTExp("")
+		}, 0, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			expSec, ok := tt.run(t)
+			if expSec != tt.wantExp || ok != tt.wantOK {
+				t.Errorf("ParseJWTExp = (%d, %t), want (%d, %t)", expSec, ok, tt.wantExp, tt.wantOK)
+			}
+		})
+	}
+}
+
 // --- ReadCodexCredential tests ---
 
 const testIDToken = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJjb2RleC10ZXN0LXN1YiIsImVtYWlsIjoidGVzdEBleGFtcGxlLmNvbSIsImV4cCI6OTk5OTk5OTk5OX0.testsig"
@@ -203,6 +259,9 @@ func TestReadCodexCredential(t *testing.T) {
 		run  func(t *testing.T)
 	}{
 		{"happy path", func(t *testing.T) {
+			// Codex Credential.ExpiresAt is always 0 — even with a valid id_token
+			// present, parseCodexCredFull does not derive expiry (the refresh gate
+			// reads the access-token JWT exp on demand in codex.StoredExpiresAt).
 			body := `{"tokens":{"access_token":"tok","refresh_token":"ref","id_token":"` + testIDToken + `"}}`
 			writeAuth(t, body)
 			cred, err := ReadCodexCredential(context.Background())
@@ -213,19 +272,19 @@ func TestReadCodexCredential(t *testing.T) {
 			if cred.RefreshToken != "ref" {
 				t.Errorf("RefreshToken = %q, want ref", cred.RefreshToken)
 			}
-			if cred.ExpiresAt != 9999999999000 {
-				t.Errorf("ExpiresAt = %d, want 9999999999000", cred.ExpiresAt)
+			if cred.ExpiresAt != 0 {
+				t.Errorf("ExpiresAt = %d, want 0 (codex never derives ExpiresAt)", cred.ExpiresAt)
 			}
 			if len(cred.Raw) == 0 {
 				t.Error("Raw should be non-empty")
 			}
 		}},
-		{"no id token", func(t *testing.T) {
+		{"no id token still zero expiry", func(t *testing.T) {
 			writeAuth(t, `{"tokens":{"access_token":"tok","refresh_token":"ref"}}`)
 			cred, err := ReadCodexCredential(context.Background())
 			testutil.WantNoErr(t, err)
 			if cred.ExpiresAt != 0 {
-				t.Errorf("ExpiresAt = %d, want 0 when id_token absent", cred.ExpiresAt)
+				t.Errorf("ExpiresAt = %d, want 0", cred.ExpiresAt)
 			}
 			if cred.AccessToken != "tok" {
 				t.Errorf("AccessToken = %q, want tok", cred.AccessToken)
@@ -264,18 +323,6 @@ func TestReadCodexCredential(t *testing.T) {
 				t.Errorf("Raw = %q, want %q", cred.Raw, body)
 			}
 		}},
-		{"malformed id token sets expires at zero", func(t *testing.T) {
-			// D4: malformed-but-present id_token sets ExpiresAt=0 without error.
-			writeAuth(t, `{"tokens":{"access_token":"tok","id_token":"not.a.valid.jwt.at.all"}}`)
-			cred, err := ReadCodexCredential(context.Background())
-			testutil.WantNoErr(t, err)
-			if cred.ExpiresAt != 0 {
-				t.Errorf("ExpiresAt = %d, want 0 for malformed id_token", cred.ExpiresAt)
-			}
-			if cred.AccessToken != "tok" {
-				t.Errorf("AccessToken = %q, want tok", cred.AccessToken)
-			}
-		}},
 		{"api key mode fails", func(t *testing.T) {
 			// D7 / A7: auth_mode=="api_key" with no tokens object → fail-closed with ErrCodexTokenNotFound.
 			writeAuth(t, `{"auth_mode":"api_key","OPENAI_API_KEY":"sk-test"}`)

internal/providers/codex/account.go

@@ -13,7 +13,6 @@ type rawCodexTokens struct {
 	Tokens struct {
 		AccessToken  string `json:"access_token"`
 		RefreshToken string `json:"refresh_token"`
-		IDToken      string `json:"id_token"`
 	} `json:"tokens"`
 }
 
@@ -43,16 +42,15 @@ func StoredRefreshToken(a accounts.Account) string {
 }
 
 // StoredExpiresAt returns the access-token expiry as milliseconds since epoch,
-// derived from the exp claim of tokens.id_token in a.RawBlob.
-// Returns 0 if id_token is absent, the JWT is malformed, or exp is absent.
+// decoded from the exp claim of tokens.access_token in a.RawBlob. Returns 0 if
+// the access token is absent or not a JWT carrying an exp claim — in which case
+// the caller performs no proactive refresh and relies on the usage call (an
+// expired-but-opaque token surfaces via the usage endpoint's 401, which maps to
+// an actionable `codex login` hint). The OIDC id_token is short-lived and is
+// NOT used here — it is identity-only (sub/email); see reconcile.go.
 func StoredExpiresAt(a accounts.Account) int64 {
-	idTok := parseStoredRaw(a).Tokens.IDToken
-	if idTok == "" {
-		return 0
+	if expSec, ok := cred.ParseJWTExp(StoredAccessToken(a)); ok {
+		return expSec * 1000
 	}
-	_, _, expSec, err := cred.ParseCodexIDToken(idTok)
-	if err != nil {
-		return 0
-	}
-	return expSec * 1000
+	return 0
 }