From 17481a5745af429e3ac73822dbf99cea13f1fcd5 Mon Sep 17 00:00:00 2001 From: zakyirsyaad Date: Sun, 5 Jul 2026 20:21:36 +0700 Subject: [PATCH 1/8] feat: add Codex access token accounts --- .../plans/2026-07-05-codex-access-token.md | 65 +++ src-tauri/Cargo.lock | 163 ++++++- src-tauri/Cargo.toml | 2 + src-tauri/src/api/usage.rs | 446 +++++++++++++++++- src-tauri/src/auth/storage.rs | 4 +- src-tauri/src/auth/switcher.rs | 138 +++++- src-tauri/src/auth/token_refresh.rs | 4 +- src-tauri/src/commands/account.rs | 39 +- src-tauri/src/commands/account_stats.rs | 4 +- src-tauri/src/commands/usage.rs | 18 +- src-tauri/src/lib.rs | 13 +- src-tauri/src/types.rs | 184 +++++++- src-tauri/src/web.rs | 26 +- src/App.tsx | 2 + src/components/AccountCard.tsx | 7 +- src/components/AddAccountModal.tsx | 84 +++- src/hooks/useAccounts.ts | 17 + src/types/index.ts | 2 +- 18 files changed, 1144 insertions(+), 74 deletions(-) create mode 100644 docs/superpowers/plans/2026-07-05-codex-access-token.md diff --git a/docs/superpowers/plans/2026-07-05-codex-access-token.md b/docs/superpowers/plans/2026-07-05-codex-access-token.md new file mode 100644 index 00000000..eeb603b1 --- /dev/null +++ b/docs/superpowers/plans/2026-07-05-codex-access-token.md @@ -0,0 +1,65 @@ +# Codex Access Token Account Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add accounts from a pasted `CODEX_ACCESS_TOKEN`, switch to them using `codex login --with-access-token`, and fetch K12/Codex usage/metadata in the same shape Codex CLI uses. + +**Architecture:** Store access-token accounts as a dedicated auth mode so existing OAuth and `auth.json` behavior stays unchanged. Parse JWT claims for display metadata, delegate actual login materialization to the Codex CLI via stdin during account switch, and use Codex's AgentIdentity request signing for access-token usage reads. + +**Tech Stack:** Rust/Tauri backend, React/TypeScript frontend, existing account store and modal patterns. + +--- + +### Task 1: Backend Auth Model + +**Files:** +- Modify: `src-tauri/src/types.rs` +- Modify: `src-tauri/src/auth/switcher.rs` + +- [x] Add a failing Rust test for `StoredAccount::new_codex_access_token`. +- [x] Add `AuthMode::CodexAccessToken` and `AuthData::CodexAccessToken`. +- [x] Parse JWT payload claims `email`, `plan_type`, and `account_id` for display metadata. +- [x] Keep existing API key and ChatGPT OAuth serialization unchanged. + +### Task 2: Backend Commands and Switch Flow + +**Files:** +- Modify: `src-tauri/src/commands/account.rs` +- Modify: `src-tauri/src/lib.rs` +- Modify: `src-tauri/src/web.rs` + +- [x] Add `add_account_from_access_token(name, access_token)`. +- [x] Validate non-empty token input. +- [x] Switch access-token accounts by spawning `codex login --with-access-token` and writing the token to stdin. +- [x] Register the command for desktop Tauri and LAN web invocation. + +### Task 3: Frontend UI + +**Files:** +- Modify: `src/hooks/useAccounts.ts` +- Modify: `src/components/AddAccountModal.tsx` +- Modify: `src/App.tsx` + +- [x] Add hook method `addFromAccessToken`. +- [x] Add an `Access Token` tab with a password textarea. +- [x] Submit to backend, reload accounts, and refresh usage as existing add-account paths do. + +### Task 4: K12/Codex Usage and Metadata + +**Files:** +- Modify: `src-tauri/src/api/usage.rs` +- Modify: `src-tauri/src/commands/usage.rs` +- Modify: `src-tauri/src/types.rs` + +- [x] Fetch access-token usage from `https://chatgpt.com/backend-api/wham/usage`. +- [x] Register/decrypt AgentIdentity task IDs when needed and send `AgentAssertion` auth headers. +- [x] Send `ChatGPT-Account-Id` from the access-token claims when available. +- [x] Accept Codex usage payload aliases such as `primary`, `secondary`, `window_duration_mins`, and `resets_at`. +- [x] Refresh stored email/plan metadata for access-token accounts from JWT claims. + +### Task 5: Verification + +- [x] Run targeted Rust tests for token parsing and Codex endpoint parsing. +- [x] Run `cargo test --manifest-path src-tauri/Cargo.toml`. +- [x] Run `pnpm build`. +- [x] Smoke-test `k12_at` through the local web endpoint: metadata returns `k12`, usage returns `error: null` with 5-hour and weekly windows. diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index bdfb8f07..c39b681f 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -250,6 +250,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64ct" +version = "1.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" + [[package]] name = "bit-set" version = "0.8.0" @@ -280,6 +286,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "blake2" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" +dependencies = [ + "digest", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -532,7 +547,9 @@ dependencies = [ "base64 0.22.1", "chacha20poly1305", "chrono", + "crypto_box", "dirs", + "ed25519-dalek", "flate2", "futures", "pbkdf2", @@ -575,6 +592,12 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "const-oid" +version = "0.9.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" + [[package]] name = "convert_case" version = "0.4.0" @@ -685,6 +708,36 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto_box" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16182b4f39a82ec8a6851155cc4c0cda3065bb1db33651726a29e1951de0f009" +dependencies = [ + "aead", + "blake2", + "crypto_secretbox", + "curve25519-dalek", + "salsa20", + "subtle", + "zeroize", +] + +[[package]] +name = "crypto_secretbox" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d6cf87adf719ddf43a805e92c6870a531aedda35ff640442cbaf8674e141e1" +dependencies = [ + "aead", + "cipher", + "generic-array", + "poly1305", + "salsa20", + "subtle", + "zeroize", +] + [[package]] name = "cssparser" version = "0.29.6" @@ -735,6 +788,33 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "curve25519-dalek" +version = "4.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be" +dependencies = [ + "cfg-if", + "cpufeatures", + "curve25519-dalek-derive", + "digest", + "fiat-crypto", + "rustc_version", + "subtle", + "zeroize", +] + +[[package]] +name = "curve25519-dalek-derive" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "darling" version = "0.23.0" @@ -769,6 +849,16 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "der" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" +dependencies = [ + "const-oid", + "zeroize", +] + [[package]] name = "deranged" version = "0.5.8" @@ -953,6 +1043,30 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" +[[package]] +name = "ed25519" +version = "2.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" +dependencies = [ + "pkcs8", + "signature", +] + +[[package]] +name = "ed25519-dalek" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" +dependencies = [ + "curve25519-dalek", + "ed25519", + "serde", + "sha2", + "subtle", + "zeroize", +] + [[package]] name = "embed-resource" version = "3.0.7" @@ -1072,6 +1186,12 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "fiat-crypto" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d" + [[package]] name = "field-offset" version = "0.3.6" @@ -1405,6 +1525,7 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", + "zeroize", ] [[package]] @@ -3018,6 +3139,16 @@ dependencies = [ "futures-io", ] +[[package]] +name = "pkcs8" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" +dependencies = [ + "der", + "spki", +] + [[package]] name = "pkg-config" version = "0.3.32" @@ -3643,6 +3774,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "salsa20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213" +dependencies = [ + "cipher", +] + [[package]] name = "same-file" version = "1.0.6" @@ -3994,6 +4134,15 @@ dependencies = [ "libc", ] +[[package]] +name = "signature" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" +dependencies = [ + "rand_core 0.6.4", +] + [[package]] name = "simd-adler32" version = "0.3.8" @@ -4098,6 +4247,16 @@ dependencies = [ "system-deps", ] +[[package]] +name = "spki" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d" +dependencies = [ + "base64ct", + "der", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -5049,9 +5208,9 @@ checksum = "bc7d623258602320d5c55d1bc22793b57daff0ec7efc270ea7d55ce1d5f5471c" [[package]] name = "typenum" -version = "1.19.0" +version = "1.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" [[package]] name = "uds_windows" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 584d0a9b..fc494d0f 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -39,3 +39,5 @@ url = "2" flate2 = "1" chacha20poly1305 = "0.10" pbkdf2 = "0.12" +ed25519-dalek = { version = "2.2", features = ["pkcs8"] } +crypto_box = { version = "0.9", features = ["seal"] } diff --git a/src-tauri/src/api/usage.rs b/src-tauri/src/api/usage.rs index 6c4c1331..90800251 100644 --- a/src-tauri/src/api/usage.rs +++ b/src-tauri/src/api/usage.rs @@ -1,35 +1,47 @@ //! Usage API client for fetching rate limits and credits use anyhow::{Context, Result}; -use chrono::{DateTime, Utc}; +use base64::Engine as _; +use chrono::{DateTime, SecondsFormat, Utc}; +use crypto_box::SecretKey as Curve25519SecretKey; +use ed25519_dalek::pkcs8::DecodePrivateKey; +use ed25519_dalek::{Signer as _, SigningKey}; use futures::{stream, StreamExt}; use reqwest::{ header::{HeaderMap, HeaderName, HeaderValue, AUTHORIZATION, USER_AGENT}, StatusCode, }; -use serde::Deserialize; +use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; -use std::collections::HashMap; +use sha2::{Digest as _, Sha512}; +use std::collections::{BTreeMap, HashMap}; use crate::auth::{ensure_chatgpt_tokens_fresh, refresh_chatgpt_tokens}; use crate::types::{ - AuthData, CreditStatusDetails, RateLimitDetails, RateLimitStatusPayload, RateLimitWindow, - StoredAccount, UsageInfo, + parse_codex_access_token_claims, AuthData, CreditStatusDetails, RateLimitDetails, + RateLimitStatusPayload, RateLimitWindow, StoredAccount, UsageInfo, }; const CHATGPT_BACKEND_API: &str = "https://chatgpt.com/backend-api"; const CHATGPT_ACCOUNTS_CHECK_API: &str = "https://chatgpt.com/backend-api/accounts/check/v4-2023-04-27"; +const CHATGPT_CODEX_USAGE_API: &str = "https://chatgpt.com/backend-api/wham/usage"; const CHATGPT_CODEX_RESPONSES_API: &str = "https://chatgpt.com/backend-api/codex/responses"; +const AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts"; const OPENAI_API: &str = "https://api.openai.com/v1"; const CODEX_USER_AGENT: &str = "codex-cli/1.0.0"; #[derive(Debug, Clone)] pub struct ChatGptAccountMetadata { + pub email: Option, pub plan_type: Option, pub subscription_expires_at: Option>, } +fn codex_usage_url() -> &'static str { + CHATGPT_CODEX_USAGE_API +} + #[derive(Debug, Deserialize)] struct AccountsCheckResponse { #[serde(default)] @@ -56,6 +68,24 @@ struct AccountsCheckEntitlement { expires_at: Option>, } +#[derive(Debug, Serialize)] +struct RegisterTaskRequest { + timestamp: String, + signature: String, +} + +#[derive(Debug, Deserialize)] +struct RegisterTaskResponse { + #[serde(default)] + task_id: Option, + #[serde(default, rename = "taskId")] + task_id_camel: Option, + #[serde(default)] + encrypted_task_id: Option, + #[serde(default, rename = "encryptedTaskId")] + encrypted_task_id_camel: Option, +} + /// Get usage information for an account pub async fn get_account_usage(account: &StoredAccount) -> Result { println!("[Usage] Fetching usage for account: {}", account.name); @@ -79,6 +109,7 @@ pub async fn get_account_usage(account: &StoredAccount) -> Result { }) } AuthData::ChatGPT { .. } => get_usage_with_chatgpt_auth(account).await, + AuthData::CodexAccessToken { .. } => get_usage_with_codex_access_token(account).await, } } @@ -92,6 +123,10 @@ pub async fn warmup_account(account: &StoredAccount) -> Result<()> { match &account.auth_data { AuthData::ApiKey { key } => warmup_with_api_key(key).await, AuthData::ChatGPT { .. } => warmup_with_chatgpt_auth(account).await, + AuthData::CodexAccessToken { .. } => { + println!("[Warmup] Skipping warm-up for CODEX_ACCESS_TOKEN account"); + Ok(()) + } } } @@ -121,6 +156,7 @@ pub async fn fetch_chatgpt_account_metadata( .context("Accounts check response did not include an account entry")?; Ok(ChatGptAccountMetadata { + email: None, plan_type: selected_entry .account .as_ref() @@ -132,6 +168,21 @@ pub async fn fetch_chatgpt_account_metadata( }) } +pub async fn fetch_codex_access_token_account_metadata( + account: &StoredAccount, +) -> Result { + let AuthData::CodexAccessToken { token, .. } = &account.auth_data else { + anyhow::bail!("Account is not using a Codex access token"); + }; + let claims = parse_codex_access_token_claims(token); + + Ok(ChatGptAccountMetadata { + email: claims.email.or_else(|| account.email.clone()), + plan_type: claims.plan_type.or_else(|| account.plan_type.clone()), + subscription_expires_at: account.subscription_expires_at, + }) +} + async fn get_usage_with_chatgpt_auth(account: &StoredAccount) -> Result { let fresh_account = ensure_chatgpt_tokens_fresh(account).await?; let (access_token, chatgpt_account_id) = extract_chatgpt_auth(&fresh_account)?; @@ -156,6 +207,13 @@ async fn get_usage_with_chatgpt_auth(account: &StoredAccount) -> Result Result { + let headers = build_codex_access_token_headers(account).await?; + let response = send_codex_usage_request(headers).await?; + + parse_usage_response(&account.id, &account.name, response).await +} + async fn parse_usage_response( account_id: &str, account_name: &str, @@ -177,10 +235,7 @@ async fn parse_usage_response( .text() .await .context("Failed to read response body")?; - println!( - "[Usage] Response body: {}", - &body_text[..body_text.len().min(200)] - ); + println!("[Usage] Response body received ({} bytes)", body_text.len()); let payload: RateLimitStatusPayload = serde_json::from_str(&body_text).context("Failed to parse usage response")?; @@ -307,6 +362,191 @@ fn build_chatgpt_headers( Ok(headers) } +async fn build_codex_access_token_headers(account: &StoredAccount) -> Result { + let AuthData::CodexAccessToken { + token, + account_id, + agent_runtime_id, + agent_private_key, + chatgpt_account_is_fedramp, + task_id, + .. + } = &account.auth_data + else { + anyhow::bail!("Account is not using a Codex access token"); + }; + + let claims = crate::types::parse_codex_access_token_claims(token); + let account_id = account_id.as_deref().or(claims.account_id.as_deref()); + let agent_runtime_id = agent_runtime_id + .as_deref() + .or(claims.agent_runtime_id.as_deref()); + let agent_private_key = agent_private_key + .as_deref() + .or(claims.agent_private_key.as_deref()); + let is_fedramp = *chatgpt_account_is_fedramp || claims.chatgpt_account_is_fedramp; + + let mut headers = HeaderMap::new(); + headers.insert(USER_AGENT, HeaderValue::from_static(CODEX_USER_AGENT)); + + if let (Some(agent_runtime_id), Some(agent_private_key)) = (agent_runtime_id, agent_private_key) + { + let task_id = match task_id.as_deref().filter(|value| !value.trim().is_empty()) { + Some(task_id) => task_id.to_string(), + None => register_agent_task(agent_runtime_id, agent_private_key).await?, + }; + let authorization = + authorization_header_for_agent_task(agent_runtime_id, agent_private_key, &task_id)?; + headers.insert( + AUTHORIZATION, + HeaderValue::from_str(&authorization).context("Invalid Agent Identity header")?, + ); + } else { + headers.insert( + AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer {}", token.trim())) + .context("Invalid access token")?, + ); + } + + if let Some(acc_id) = account_id { + println!("[Usage] Using Codex Account ID: {acc_id}"); + if let Ok(header_name) = HeaderName::from_bytes(b"ChatGPT-Account-ID") { + if let Ok(header_value) = HeaderValue::from_str(acc_id) { + headers.insert(header_name, header_value); + } + } + } + + if is_fedramp { + headers.insert( + HeaderName::from_static("x-openai-fedramp"), + HeaderValue::from_static("true"), + ); + } + + Ok(headers) +} + +async fn register_agent_task(agent_runtime_id: &str, agent_private_key: &str) -> Result { + let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + let request = RegisterTaskRequest { + timestamp: timestamp.clone(), + signature: sign_task_registration_payload(agent_runtime_id, agent_private_key, ×tamp)?, + }; + let url = format!( + "{}/v1/agent/{}/task/register", + AGENT_IDENTITY_AUTHAPI_BASE_URL.trim_end_matches('/'), + agent_runtime_id + ); + + let response = reqwest::Client::new() + .post(&url) + .json(&request) + .send() + .await + .with_context(|| format!("Failed to register Agent Identity task at {url}"))?; + + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + anyhow::bail!("Agent Identity task registration failed: {status} - {body}"); + } + + let payload: RegisterTaskResponse = response + .json() + .await + .context("Failed to parse Agent Identity task registration response")?; + + task_id_from_register_task_response(agent_private_key, payload) +} + +fn task_id_from_register_task_response( + private_key_pkcs8_base64: &str, + response: RegisterTaskResponse, +) -> Result { + if let Some(task_id) = response.task_id.or(response.task_id_camel) { + return Ok(task_id); + } + + let encrypted_task_id = response + .encrypted_task_id + .or(response.encrypted_task_id_camel) + .context("Agent Identity task registration response omitted task_id")?; + + decrypt_task_id_response(private_key_pkcs8_base64, &encrypted_task_id) +} + +fn sign_task_registration_payload( + agent_runtime_id: &str, + private_key_pkcs8_base64: &str, + timestamp: &str, +) -> Result { + sign_agent_identity_payload( + private_key_pkcs8_base64, + &format!("{agent_runtime_id}:{timestamp}"), + ) +} + +fn authorization_header_for_agent_task( + agent_runtime_id: &str, + private_key_pkcs8_base64: &str, + task_id: &str, +) -> Result { + let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + let signature = sign_agent_identity_payload( + private_key_pkcs8_base64, + &format!("{agent_runtime_id}:{task_id}:{timestamp}"), + )?; + let payload = serde_json::to_vec(&BTreeMap::from([ + ("agent_runtime_id", agent_runtime_id), + ("signature", signature.as_str()), + ("task_id", task_id), + ("timestamp", timestamp.as_str()), + ])) + .context("Failed to serialize Agent Identity assertion")?; + let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(payload); + Ok(format!("AgentAssertion {encoded}")) +} + +fn sign_agent_identity_payload(private_key_pkcs8_base64: &str, payload: &str) -> Result { + let signing_key = signing_key_from_private_key_pkcs8_base64(private_key_pkcs8_base64)?; + let signature = signing_key.sign(payload.as_bytes()); + Ok(base64::engine::general_purpose::STANDARD.encode(signature.to_bytes())) +} + +fn decrypt_task_id_response( + private_key_pkcs8_base64: &str, + encrypted_task_id: &str, +) -> Result { + let signing_key = signing_key_from_private_key_pkcs8_base64(private_key_pkcs8_base64)?; + let ciphertext = base64::engine::general_purpose::STANDARD + .decode(encrypted_task_id) + .context("Encrypted Agent Identity task id is not valid base64")?; + let plaintext = curve25519_secret_key_from_signing_key(&signing_key) + .unseal(&ciphertext) + .map_err(|_| anyhow::anyhow!("Failed to decrypt Agent Identity task id"))?; + String::from_utf8(plaintext).context("Decrypted Agent Identity task id is not valid UTF-8") +} + +fn signing_key_from_private_key_pkcs8_base64(private_key_pkcs8_base64: &str) -> Result { + let private_key = base64::engine::general_purpose::STANDARD + .decode(private_key_pkcs8_base64) + .context("Agent Identity private key is not valid base64")?; + SigningKey::from_pkcs8_der(&private_key) + .context("Agent Identity private key is not valid PKCS#8") +} + +fn curve25519_secret_key_from_signing_key(signing_key: &SigningKey) -> Curve25519SecretKey { + let digest = Sha512::digest(signing_key.to_bytes()); + let mut secret_key = [0u8; 32]; + secret_key.copy_from_slice(&digest[..32]); + secret_key[0] &= 248; + secret_key[31] &= 127; + secret_key[31] |= 64; + Curve25519SecretKey::from_bytes(secret_key) +} + fn extract_chatgpt_auth(account: &StoredAccount) -> Result<(&str, Option<&str>)> { match &account.auth_data { AuthData::ChatGPT { @@ -314,6 +554,9 @@ fn extract_chatgpt_auth(account: &StoredAccount) -> Result<(&str, Option<&str>)> account_id, .. } => Ok((access_token.as_str(), account_id.as_deref())), + AuthData::CodexAccessToken { + token, account_id, .. + } => Ok((token.as_str(), account_id.as_deref())), AuthData::ApiKey { .. } => anyhow::bail!("Account is not using ChatGPT OAuth"), } } @@ -330,16 +573,23 @@ async fn send_chatgpt_usage_request( .await } +async fn send_codex_usage_request(headers: HeaderMap) -> Result { + send_get_request_with_headers(codex_usage_url(), headers).await +} + async fn send_chatgpt_get_request( url: &str, access_token: &str, chatgpt_account_id: Option<&str>, ) -> Result { - let client = reqwest::Client::new(); let headers = build_chatgpt_headers(access_token, chatgpt_account_id)?; + send_get_request_with_headers(url, headers).await +} + +async fn send_get_request_with_headers(url: &str, headers: HeaderMap) -> Result { println!("[Usage] Requesting: {url}"); - client + reqwest::Client::new() .get(url) .headers(headers) .send() @@ -458,16 +708,10 @@ fn convert_payload_to_usage_info(account_id: &str, payload: RateLimitStatusPaylo account_id: account_id.to_string(), plan_type: Some(payload.plan_type), primary_used_percent: primary.as_ref().map(|w| w.used_percent), - primary_window_minutes: primary - .as_ref() - .and_then(|w| w.limit_window_seconds) - .map(|s| (i64::from(s) + 59) / 60), + primary_window_minutes: primary.as_ref().and_then(window_minutes), primary_resets_at: primary.as_ref().and_then(|w| w.reset_at), secondary_used_percent: secondary.as_ref().map(|w| w.used_percent), - secondary_window_minutes: secondary - .as_ref() - .and_then(|w| w.limit_window_seconds) - .map(|s| (i64::from(s) + 59) / 60), + secondary_window_minutes: secondary.as_ref().and_then(window_minutes), secondary_resets_at: secondary.as_ref().and_then(|w| w.reset_at), has_credits: credits.as_ref().map(|c| c.has_credits), unlimited_credits: credits.as_ref().map(|c| c.unlimited), @@ -476,6 +720,14 @@ fn convert_payload_to_usage_info(account_id: &str, payload: RateLimitStatusPaylo } } +fn window_minutes(window: &RateLimitWindow) -> Option { + window.window_duration_mins.or_else(|| { + window + .limit_window_seconds + .map(|s| (i64::from(s) + 59) / 60) + }) +} + fn extract_rate_limits( rate_limit: Option, ) -> (Option, Option) { @@ -511,3 +763,159 @@ pub async fn refresh_all_usage(accounts: &[StoredAccount]) -> Vec { println!("[Usage] Refresh complete"); results } + +#[cfg(test)] +mod tests { + use super::{ + authorization_header_for_agent_task, build_codex_access_token_headers, codex_usage_url, + convert_payload_to_usage_info, + }; + use crate::types::{AuthData, AuthMode, RateLimitStatusPayload, StoredAccount}; + use base64::Engine as _; + use chrono::Utc; + use ed25519_dalek::pkcs8::EncodePrivateKey; + use ed25519_dalek::SigningKey; + use reqwest::header::AUTHORIZATION; + use serde_json::json; + use uuid::Uuid; + + #[test] + fn codex_access_token_usage_uses_chatgpt_wham_usage_endpoint() { + assert_eq!( + codex_usage_url(), + "https://chatgpt.com/backend-api/wham/usage" + ); + } + + #[test] + fn codex_access_token_headers_include_account_id_when_available() { + let headers = super::build_chatgpt_headers("token", Some("acc_123")).unwrap(); + + assert_eq!( + headers + .get("ChatGPT-Account-Id") + .and_then(|value| value.to_str().ok()), + Some("acc_123") + ); + } + + #[test] + fn agent_identity_authorization_header_uses_agent_assertion_scheme() { + let signing_key = SigningKey::from_bytes(&[7u8; 32]); + let private_key = signing_key.to_pkcs8_der().unwrap(); + let private_key_base64 = + base64::engine::general_purpose::STANDARD.encode(private_key.as_bytes()); + + let header = authorization_header_for_agent_task( + "agent-runtime-123", + &private_key_base64, + "task-123", + ) + .unwrap(); + let encoded = header.strip_prefix("AgentAssertion ").unwrap(); + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(encoded) + .unwrap(); + let envelope: serde_json::Value = serde_json::from_slice(&payload).unwrap(); + + assert_eq!(envelope["agent_runtime_id"], "agent-runtime-123"); + assert_eq!(envelope["task_id"], "task-123"); + assert!(envelope["timestamp"].as_str().is_some()); + assert!(envelope["signature"].as_str().is_some()); + } + + #[tokio::test] + async fn codex_access_token_headers_use_agent_assertion_when_agent_claims_exist() { + let signing_key = SigningKey::from_bytes(&[8u8; 32]); + let private_key = signing_key.to_pkcs8_der().unwrap(); + let private_key_base64 = + base64::engine::general_purpose::STANDARD.encode(private_key.as_bytes()); + let sample_access_token = ["header", "payload", "signature"].join("."); + let account = StoredAccount { + id: Uuid::new_v4().to_string(), + name: "agent".to_string(), + email: Some("agent@example.com".to_string()), + plan_type: Some("k12".to_string()), + subscription_expires_at: None, + auth_mode: AuthMode::CodexAccessToken, + auth_data: AuthData::CodexAccessToken { + token: sample_access_token, + account_id: Some("account-123".to_string()), + agent_runtime_id: Some("agent-runtime-123".to_string()), + agent_private_key: Some(private_key_base64), + chatgpt_user_id: Some("user-123".to_string()), + chatgpt_account_is_fedramp: false, + task_id: Some("task-123".to_string()), + }, + created_at: Utc::now(), + last_used_at: None, + }; + + let headers = build_codex_access_token_headers(&account).await.unwrap(); + + assert_eq!( + headers + .get("ChatGPT-Account-ID") + .and_then(|value| value.to_str().ok()), + Some("account-123") + ); + assert!(headers + .get(AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.starts_with("AgentAssertion "))); + } + + #[tokio::test] + async fn codex_access_token_metadata_comes_from_token_claims() { + let payload = r#"{"email":"k12@example.com","plan_type":"k12","account_id":"account-123","agent_runtime_id":"agent-runtime-123"}"#; + let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(payload); + let account = + StoredAccount::new_codex_access_token("agent".to_string(), format!("h.{encoded}.s")); + + let metadata = super::fetch_codex_access_token_account_metadata(&account) + .await + .unwrap(); + + assert_eq!(metadata.email.as_deref(), Some("k12@example.com")); + assert_eq!(metadata.plan_type.as_deref(), Some("k12")); + } + + #[test] + fn converts_codex_usage_payload_aliases() { + let payload: RateLimitStatusPayload = serde_json::from_value(json!({ + "plan_type": "k12", + "rate_limit": { + "primary": { + "used_percent": 12.5, + "window_duration_mins": 300, + "resets_at": 1783250000 + }, + "secondary": { + "used_percent": 2.25, + "window_duration_mins": 10080, + "resets_at": 1783850000 + } + }, + "credits": { + "has_credits": true, + "unlimited": false, + "balance": "$10.00" + } + })) + .unwrap(); + + let usage = convert_payload_to_usage_info("acc_123", payload); + + assert_eq!(usage.plan_type.as_deref(), Some("k12")); + assert_eq!(usage.primary_used_percent, Some(12.5)); + assert_eq!(usage.primary_window_minutes, Some(300)); + assert_eq!(usage.primary_resets_at, Some(1783250000)); + assert_eq!(usage.secondary_used_percent, Some(2.25)); + assert_eq!(usage.secondary_window_minutes, Some(10080)); + assert_eq!(usage.secondary_resets_at, Some(1783850000)); + assert_eq!(usage.has_credits, Some(true)); + assert_eq!(usage.unlimited_credits, Some(false)); + assert_eq!(usage.credits_balance.as_deref(), Some("$10.00")); + assert_eq!(usage.error, None); + } +} diff --git a/src-tauri/src/auth/storage.rs b/src-tauri/src/auth/storage.rs index ace0f25f..54514851 100644 --- a/src-tauri/src/auth/storage.rs +++ b/src-tauri/src/auth/storage.rs @@ -270,8 +270,8 @@ pub fn update_account_chatgpt_tokens( *stored_account_id = Some(new_account_id); } } - AuthData::ApiKey { .. } => { - anyhow::bail!("Cannot update OAuth tokens for an API key account"); + AuthData::ApiKey { .. } | AuthData::CodexAccessToken { .. } => { + anyhow::bail!("Cannot update OAuth tokens for this account type"); } } diff --git a/src-tauri/src/auth/switcher.rs b/src-tauri/src/auth/switcher.rs index e213574e..94f4b3d0 100644 --- a/src-tauri/src/auth/switcher.rs +++ b/src-tauri/src/auth/switcher.rs @@ -1,7 +1,9 @@ //! Account switching logic - writes credentials to ~/.codex/auth.json use std::fs; +use std::io::Write; use std::path::PathBuf; +use std::process::{Command, Stdio}; use anyhow::{Context, Result}; use chrono::Utc; @@ -28,6 +30,10 @@ pub fn get_codex_auth_file() -> Result { /// Switch to a specific account by writing its credentials to ~/.codex/auth.json pub fn switch_to_account(account: &StoredAccount) -> Result<()> { + if let AuthData::CodexAccessToken { token, .. } = &account.auth_data { + return login_with_codex_access_token(token); + } + let codex_home = get_codex_home()?; // Ensure the codex home directory exists @@ -58,9 +64,12 @@ pub fn switch_to_account(account: &StoredAccount) -> Result<()> { fn create_auth_json(account: &StoredAccount) -> Result { match &account.auth_data { AuthData::ApiKey { key } => Ok(AuthDotJson { + auth_mode: None, openai_api_key: Some(key.clone()), tokens: None, last_refresh: None, + agent_identity: None, + personal_access_token: None, }), AuthData::ChatGPT { id_token, @@ -68,6 +77,7 @@ fn create_auth_json(account: &StoredAccount) -> Result { refresh_token, account_id, } => Ok(AuthDotJson { + auth_mode: None, openai_api_key: None, tokens: Some(TokenData { id_token: id_token.clone(), @@ -76,7 +86,116 @@ fn create_auth_json(account: &StoredAccount) -> Result { account_id: account_id.clone(), }), last_refresh: Some(Utc::now()), + agent_identity: None, + personal_access_token: None, }), + AuthData::CodexAccessToken { token, .. } => Ok(create_access_token_auth_json(token)), + } +} + +fn create_access_token_auth_json(token: &str) -> AuthDotJson { + let trimmed = token.trim().to_string(); + if trimmed.starts_with("at-") { + AuthDotJson { + auth_mode: None, + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: Some(trimmed), + } + } else { + AuthDotJson { + auth_mode: Some("agentIdentity".to_string()), + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: Some(serde_json::Value::String(trimmed)), + personal_access_token: None, + } + } +} + +fn login_with_codex_access_token(token: &str) -> Result<()> { + let trimmed = token.trim(); + if trimmed.is_empty() { + anyhow::bail!("Codex access token is empty"); + } + + let mut child = Command::new("codex") + .args(["login", "--with-access-token"]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .context("Failed to start Codex CLI. Make sure `codex` is installed and on PATH")?; + + { + let mut stdin = child + .stdin + .take() + .context("Failed to open Codex CLI stdin")?; + writeln!(stdin, "{trimmed}").context("Failed to send access token to Codex CLI")?; + } + + let output = child + .wait_with_output() + .context("Failed to wait for Codex CLI login")?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); + let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string(); + let detail = if !stderr.is_empty() { + stderr + } else if !stdout.is_empty() { + stdout + } else { + format!("exit status {}", output.status) + }; + let detail = redact_access_token_from_output(&detail, trimmed); + anyhow::bail!("Codex access token login failed: {detail}"); + } + + Ok(()) +} + +fn redact_access_token_from_output(output: &str, token: &str) -> String { + if token.is_empty() { + return output.to_string(); + } + + output.replace(token, "[redacted access token]") +} + +#[cfg(test)] +mod tests { + use super::{create_access_token_auth_json, redact_access_token_from_output}; + + #[test] + fn redacts_access_token_from_cli_error_output() { + let marker = ["sample", "token", "value"].join("-"); + let output = format!("login failed for {marker}"); + + assert_eq!( + redact_access_token_from_output(&output, &marker), + "login failed for [redacted access token]" + ); + } + + #[test] + fn creates_agent_identity_auth_json_for_codex_access_token_jwt() { + let sample_access_token = ["header", "payload", "signature"].join("."); + let auth = create_access_token_auth_json(&sample_access_token); + + assert_eq!(auth.auth_mode.as_deref(), Some("agentIdentity")); + assert_eq!( + auth.agent_identity + .as_ref() + .and_then(|value| value.as_str()), + Some(sample_access_token.as_str()) + ); + assert!(auth.tokens.is_none()); + assert!(auth.personal_access_token.is_none()); } } @@ -113,8 +232,20 @@ pub fn import_from_auth_json_contents( tokens.refresh_token, claims.account_id.or(tokens.account_id), )) + } else if let Some(agent_identity) = auth.agent_identity { + let token = match agent_identity { + serde_json::Value::String(token) => token, + _ => anyhow::bail!("auth.json agent_identity has an unsupported shape"), + }; + + Ok(StoredAccount::new_codex_access_token(account_name, token)) + } else if let Some(personal_access_token) = auth.personal_access_token { + Ok(StoredAccount::new_codex_access_token( + account_name, + personal_access_token, + )) } else { - anyhow::bail!("auth.json contains neither API key nor tokens"); + anyhow::bail!("auth.json contains neither API key, tokens, nor access-token auth"); } } @@ -138,7 +269,10 @@ pub fn read_current_auth() -> Result> { /// Check if there is an active Codex login pub fn has_active_login() -> Result { match read_current_auth()? { - Some(auth) => Ok(auth.openai_api_key.is_some() || auth.tokens.is_some()), + Some(auth) => Ok(auth.openai_api_key.is_some() + || auth.tokens.is_some() + || auth.agent_identity.is_some() + || auth.personal_access_token.is_some()), None => Ok(false), } } diff --git a/src-tauri/src/auth/token_refresh.rs b/src-tauri/src/auth/token_refresh.rs index fbba025a..c6b91842 100644 --- a/src-tauri/src/auth/token_refresh.rs +++ b/src-tauri/src/auth/token_refresh.rs @@ -25,7 +25,7 @@ struct RefreshTokenResponse { /// Returns an updated account when a refresh was performed. pub async fn ensure_chatgpt_tokens_fresh(account: &StoredAccount) -> Result { match &account.auth_data { - AuthData::ApiKey { .. } => Ok(account.clone()), + AuthData::ApiKey { .. } | AuthData::CodexAccessToken { .. } => Ok(account.clone()), AuthData::ChatGPT { access_token, .. } => { if token_expired_or_near_expiry(access_token) { println!( @@ -43,7 +43,7 @@ pub async fn ensure_chatgpt_tokens_fresh(account: &StoredAccount) -> Result Result { let (current_id_token, current_refresh_token, current_account_id) = match &account.auth_data { - AuthData::ApiKey { .. } => return Ok(account.clone()), + AuthData::ApiKey { .. } | AuthData::CodexAccessToken { .. } => return Ok(account.clone()), AuthData::ChatGPT { id_token, refresh_token, diff --git a/src-tauri/src/commands/account.rs b/src-tauri/src/commands/account.rs index 915bb77b..5bc2e6e1 100644 --- a/src-tauri/src/commands/account.rs +++ b/src-tauri/src/commands/account.rs @@ -125,6 +125,32 @@ pub async fn add_account_from_auth_json_text( Ok(AccountInfo::from_stored(&stored, active_id)) } +/// Add an account from a CODEX_ACCESS_TOKEN value. +#[tauri::command] +pub async fn add_account_from_access_token( + name: String, + access_token: String, +) -> Result { + let trimmed_name = name.trim(); + if trimmed_name.is_empty() { + return Err("Account name is required".to_string()); + } + + let trimmed_token = access_token.trim(); + if trimmed_token.is_empty() { + return Err("Access token is required".to_string()); + } + + let account = + StoredAccount::new_codex_access_token(trimmed_name.to_string(), trimmed_token.to_string()); + let stored = add_account(account).map_err(|e| e.to_string())?; + + let store = load_accounts().map_err(|e| e.to_string())?; + let active_id = store.active_account_id.as_deref(); + + Ok(AccountInfo::from_stored(&stored, active_id)) +} + /// Switch to a different account #[tauri::command] pub async fn switch_account(account_id: String) -> Result<(), String> { @@ -349,20 +375,23 @@ fn encode_slim_payload_from_store(store: &AccountsStore) -> anyhow::Result SlimAccountPayload { + AuthData::ApiKey { key } => Ok(SlimAccountPayload { name: account.name.clone(), auth_type: SLIM_AUTH_API_KEY, api_key: Some(key.clone()), refresh_token: None, - }, - AuthData::ChatGPT { refresh_token, .. } => SlimAccountPayload { + }), + AuthData::ChatGPT { refresh_token, .. } => Ok(SlimAccountPayload { name: account.name.clone(), auth_type: SLIM_AUTH_CHATGPT, api_key: None, refresh_token: Some(refresh_token.clone()), - }, + }), + AuthData::CodexAccessToken { .. } => anyhow::bail!( + "Slim export does not support CODEX_ACCESS_TOKEN accounts. Use full encrypted export instead." + ), }) - .collect(); + .collect::>>()?; let payload = SlimPayload { version: SLIM_FORMAT_VERSION, diff --git a/src-tauri/src/commands/account_stats.rs b/src-tauri/src/commands/account_stats.rs index 0bb1273b..d18b7ac5 100644 --- a/src-tauri/src/commands/account_stats.rs +++ b/src-tauri/src/commands/account_stats.rs @@ -445,7 +445,9 @@ fn extract_chatgpt_auth(account: &StoredAccount) -> anyhow::Result<(&str, Option account_id, .. } => Ok((access_token.as_str(), account_id.as_deref())), - AuthData::ApiKey { .. } => anyhow::bail!("Account is not using ChatGPT OAuth"), + AuthData::ApiKey { .. } | AuthData::CodexAccessToken { .. } => { + anyhow::bail!("Account is not using ChatGPT OAuth") + } } } diff --git a/src-tauri/src/commands/usage.rs b/src-tauri/src/commands/usage.rs index 045a8e5d..11a8cb9b 100644 --- a/src-tauri/src/commands/usage.rs +++ b/src-tauri/src/commands/usage.rs @@ -1,8 +1,8 @@ //! Usage query Tauri commands use crate::api::usage::{ - fetch_chatgpt_account_metadata, get_account_usage, refresh_all_usage, - warmup_account as send_warmup, + fetch_chatgpt_account_metadata, fetch_codex_access_token_account_metadata, get_account_usage, + refresh_all_usage, warmup_account as send_warmup, }; use crate::auth::{get_account, load_accounts, refresh_chatgpt_tokens, update_account_metadata}; use crate::types::{AccountInfo, AuthData, UsageInfo, WarmupSummary}; @@ -53,7 +53,21 @@ pub async fn refresh_account_metadata(account_id: String) -> Result { + let live_metadata = fetch_codex_access_token_account_metadata(&account) + .await + .map_err(|e| e.to_string())?; + + update_account_metadata( + &account_id, None, + live_metadata.email, live_metadata.plan_type, Some(live_metadata.subscription_expires_at), ) diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 72956a7c..768ae710 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -11,12 +11,12 @@ pub mod types; pub mod web; use commands::{ - ack_close_behavior_prompt, add_account_from_file, cancel_login, check_codex_processes, - complete_close_behavior, complete_login, delete_account, export_accounts_full_encrypted_file, - export_accounts_slim_text, get_account_usage_stats, get_active_account_info, - get_dock_display_mode, get_masked_account_ids, get_usage, hide_tray_window, - import_accounts_full_encrypted_file, import_accounts_slim_text, kill_codex_processes, - list_accounts, open_main_window, quit_app, refresh_account_metadata, + ack_close_behavior_prompt, add_account_from_access_token, add_account_from_file, cancel_login, + check_codex_processes, complete_close_behavior, complete_login, delete_account, + export_accounts_full_encrypted_file, export_accounts_slim_text, get_account_usage_stats, + get_active_account_info, get_dock_display_mode, get_masked_account_ids, get_usage, + hide_tray_window, import_accounts_full_encrypted_file, import_accounts_slim_text, + kill_codex_processes, list_accounts, open_main_window, quit_app, refresh_account_metadata, refresh_all_accounts_usage, rename_account, report_usage, set_dock_display_mode, set_masked_account_ids, start_login, switch_account, warmup_account, warmup_all_accounts, }; @@ -65,6 +65,7 @@ pub fn run() { list_accounts, get_active_account_info, add_account_from_file, + add_account_from_access_token, switch_account, delete_account, rename_account, diff --git a/src-tauri/src/types.rs b/src-tauri/src/types.rs index de1abfd2..81c662db 100644 --- a/src-tauri/src/types.rs +++ b/src-tauri/src/types.rs @@ -138,6 +138,31 @@ impl StoredAccount { last_used_at: None, } } + + /// Create a new account from a Codex CLI access token. + pub fn new_codex_access_token(name: String, token: String) -> Self { + let claims = parse_codex_access_token_claims(&token); + + Self { + id: Uuid::new_v4().to_string(), + name, + email: claims.email.clone(), + plan_type: claims.plan_type.clone(), + subscription_expires_at: None, + auth_mode: AuthMode::CodexAccessToken, + auth_data: AuthData::CodexAccessToken { + token, + account_id: claims.account_id, + agent_runtime_id: claims.agent_runtime_id, + agent_private_key: claims.agent_private_key, + chatgpt_user_id: claims.chatgpt_user_id, + chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp, + task_id: None, + }, + created_at: Utc::now(), + last_used_at: None, + } + } } /// Authentication mode @@ -148,6 +173,8 @@ pub enum AuthMode { ApiKey, /// Using ChatGPT OAuth tokens ChatGPT, + /// Using a Codex CLI access token + CodexAccessToken, } /// Authentication data (credentials) @@ -170,6 +197,29 @@ pub enum AuthData { /// ChatGPT account ID account_id: Option, }, + /// Codex CLI access token authentication + CodexAccessToken { + /// The raw CODEX_ACCESS_TOKEN value + token: String, + /// Account ID from token claims, when present + #[serde(default)] + account_id: Option, + /// Agent runtime ID from agent identity JWT claims, when present + #[serde(default)] + agent_runtime_id: Option, + /// Agent private key from agent identity JWT claims, when present + #[serde(default)] + agent_private_key: Option, + /// ChatGPT user ID from agent identity JWT claims, when present + #[serde(default)] + chatgpt_user_id: Option, + /// Whether this account should route through FedRAMP headers + #[serde(default)] + chatgpt_account_is_fedramp: bool, + /// Optional registered task ID for signed Agent Identity requests + #[serde(default)] + task_id: Option, + }, } #[derive(Debug, Clone, Default)] @@ -180,20 +230,54 @@ pub struct ChatGptIdTokenClaims { pub subscription_expires_at: Option>, } -pub fn parse_chatgpt_id_token_claims(id_token: &str) -> ChatGptIdTokenClaims { - let parts: Vec<&str> = id_token.split('.').collect(); - if parts.len() != 3 { - return ChatGptIdTokenClaims::default(); - } +#[derive(Debug, Clone, Default)] +pub struct CodexAccessTokenClaims { + pub email: Option, + pub plan_type: Option, + pub account_id: Option, + pub agent_runtime_id: Option, + pub agent_private_key: Option, + pub chatgpt_user_id: Option, + pub chatgpt_account_is_fedramp: bool, +} - let payload = match base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(parts[1]) { - Ok(bytes) => bytes, - Err(_) => return ChatGptIdTokenClaims::default(), +pub fn parse_codex_access_token_claims(token: &str) -> CodexAccessTokenClaims { + let Some(json) = parse_jwt_payload(token) else { + return CodexAccessTokenClaims::default(); }; - let json: serde_json::Value = match serde_json::from_slice(&payload) { - Ok(value) => value, - Err(_) => return ChatGptIdTokenClaims::default(), + CodexAccessTokenClaims { + email: json.get("email").and_then(|v| v.as_str()).map(String::from), + plan_type: json + .get("plan_type") + .and_then(|v| v.as_str()) + .map(String::from), + account_id: json + .get("account_id") + .and_then(|v| v.as_str()) + .map(String::from), + agent_runtime_id: json + .get("agent_runtime_id") + .and_then(|v| v.as_str()) + .map(String::from), + agent_private_key: json + .get("agent_private_key") + .and_then(|v| v.as_str()) + .map(String::from), + chatgpt_user_id: json + .get("chatgpt_user_id") + .and_then(|v| v.as_str()) + .map(String::from), + chatgpt_account_is_fedramp: json + .get("chatgpt_account_is_fedramp") + .and_then(|v| v.as_bool()) + .unwrap_or(false), + } +} + +pub fn parse_chatgpt_id_token_claims(id_token: &str) -> ChatGptIdTokenClaims { + let Some(json) = parse_jwt_payload(id_token) else { + return ChatGptIdTokenClaims::default(); }; let auth_claims = json.get("https://api.openai.com/auth"); @@ -216,6 +300,20 @@ pub fn parse_chatgpt_id_token_claims(id_token: &str) -> ChatGptIdTokenClaims { } } +fn parse_jwt_payload(token: &str) -> Option { + let parts: Vec<&str> = token.split('.').collect(); + if parts.len() != 3 { + return None; + } + + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(parts[1]) + .or_else(|_| base64::engine::general_purpose::URL_SAFE.decode(parts[1])) + .ok()?; + + serde_json::from_slice(&payload).ok() +} + // ============================================================================ // Types for Codex's auth.json format (for compatibility) // ============================================================================ @@ -223,6 +321,9 @@ pub fn parse_chatgpt_id_token_claims(id_token: &str) -> ChatGptIdTokenClaims { /// The official Codex auth.json format #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AuthDotJson { + /// Explicit auth mode used by newer Codex auth.json formats + #[serde(default, skip_serializing_if = "Option::is_none")] + pub auth_mode: Option, /// OpenAI API key (for API key auth mode) #[serde(rename = "OPENAI_API_KEY", skip_serializing_if = "Option::is_none")] pub openai_api_key: Option, @@ -232,6 +333,12 @@ pub struct AuthDotJson { /// Last token refresh timestamp #[serde(skip_serializing_if = "Option::is_none")] pub last_refresh: Option>, + /// Agent identity auth material used by CODEX_ACCESS_TOKEN JWTs + #[serde(default, skip_serializing_if = "Option::is_none")] + pub agent_identity: Option, + /// Personal access token auth material used by access tokens prefixed with at- + #[serde(default, skip_serializing_if = "Option::is_none")] + pub personal_access_token: Option, } /// Token data stored in auth.json @@ -272,7 +379,7 @@ impl AccountInfo { AuthData::ChatGPT { id_token, .. } => { parse_chatgpt_id_token_claims(id_token).subscription_expires_at } - AuthData::ApiKey { .. } => None, + AuthData::ApiKey { .. } | AuthData::CodexAccessToken { .. } => None, }; Self { @@ -387,14 +494,20 @@ pub struct RateLimitStatusPayload { #[derive(Debug, Clone, Deserialize)] pub struct RateLimitDetails { + #[serde(alias = "primary")] pub primary_window: Option, + #[serde(alias = "secondary")] pub secondary_window: Option, } #[derive(Debug, Clone, Deserialize)] pub struct RateLimitWindow { pub used_percent: f64, + #[serde(default)] pub limit_window_seconds: Option, + #[serde(default)] + pub window_duration_mins: Option, + #[serde(alias = "resets_at")] pub reset_at: Option, } @@ -408,7 +521,10 @@ pub struct CreditStatusDetails { #[cfg(test)] mod tests { - use super::{parse_chatgpt_id_token_claims, AppSettings, DockDisplayMode, TrayDisplayMode}; + use super::{ + parse_chatgpt_id_token_claims, AppSettings, AuthData, AuthMode, DockDisplayMode, + StoredAccount, TrayDisplayMode, + }; use base64::Engine; #[test] @@ -439,4 +555,46 @@ mod tests { assert_eq!(settings.dock_display_mode, DockDisplayMode::ShowInDock); assert!(settings.close_behavior_prompt_enabled); } + + #[test] + fn creates_codex_access_token_account_from_jwt_claims() { + let payload = r#"{"account_id":"acc_env_123","agent_runtime_id":"agent-runtime-123","agent_private_key":"private-key","chatgpt_user_id":"user-123","chatgpt_account_is_fedramp":true,"email":"token-user@example.com","plan_type":"team","exp":2098606399}"#; + let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(payload); + let token = format!("header.{encoded}.signature"); + + let account = + StoredAccount::new_codex_access_token("Env Account".to_string(), token.clone()); + + assert_eq!(account.name, "Env Account"); + assert_eq!(account.email.as_deref(), Some("token-user@example.com")); + assert_eq!(account.plan_type.as_deref(), Some("team")); + assert_eq!(account.auth_mode, AuthMode::CodexAccessToken); + match account.auth_data { + AuthData::CodexAccessToken { + token: stored_token, + account_id, + agent_runtime_id, + agent_private_key, + chatgpt_user_id, + chatgpt_account_is_fedramp, + task_id, + } => { + assert_eq!(stored_token, token); + assert_eq!(account_id.as_deref(), Some("acc_env_123")); + assert_eq!(agent_runtime_id.as_deref(), Some("agent-runtime-123")); + assert_eq!(agent_private_key.as_deref(), Some("private-key")); + assert_eq!(chatgpt_user_id.as_deref(), Some("user-123")); + assert!(chatgpt_account_is_fedramp); + assert_eq!(task_id, None); + } + other => panic!("expected CodexAccessToken auth data, got {other:?}"), + } + } + + #[test] + fn serializes_codex_access_token_auth_mode_as_snake_case() { + let encoded = serde_json::to_string(&AuthMode::CodexAccessToken).unwrap(); + + assert_eq!(encoded, "\"codex_access_token\""); + } } diff --git a/src-tauri/src/web.rs b/src-tauri/src/web.rs index 1d91290e..344d6ae3 100644 --- a/src-tauri/src/web.rs +++ b/src-tauri/src/web.rs @@ -10,13 +10,13 @@ use tiny_http::{Header, Method, Request, Response, Server, StatusCode}; use tokio::runtime::Runtime; use crate::commands::{ - add_account_from_auth_json_text, add_account_from_file, cancel_login, check_codex_processes, - complete_login, delete_account, export_accounts_full_encrypted_bytes, - export_accounts_slim_text, fetch_usage, get_account_usage_stats, get_active_account_info, - get_masked_account_ids, import_accounts_full_encrypted_bytes, import_accounts_slim_text, - kill_codex_processes, list_accounts, refresh_account_metadata, refresh_all_accounts_usage, - rename_account, set_masked_account_ids, start_login, switch_account, warmup_account, - warmup_all_accounts, + add_account_from_access_token, add_account_from_auth_json_text, add_account_from_file, + cancel_login, check_codex_processes, complete_login, delete_account, + export_accounts_full_encrypted_bytes, export_accounts_slim_text, fetch_usage, + get_account_usage_stats, get_active_account_info, get_masked_account_ids, + import_accounts_full_encrypted_bytes, import_accounts_slim_text, kill_codex_processes, + list_accounts, refresh_account_metadata, refresh_all_accounts_usage, rename_account, + set_masked_account_ids, start_login, switch_account, warmup_account, warmup_all_accounts, }; #[derive(Debug, Deserialize)] @@ -58,6 +58,14 @@ struct UploadAuthJsonArgs { contents: String, } +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct AccessTokenAccountArgs { + name: String, + #[serde(alias = "access_token")] + access_token: String, +} + #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct UploadEncryptedArgs { @@ -138,6 +146,10 @@ async fn invoke_web_command(command: &str, payload: Value) -> Result { + let args: AccessTokenAccountArgs = parse_args(payload)?; + to_json(add_account_from_access_token(args.name, args.access_token).await?) + } "get_usage" => { let args: AccountIdArgs = parse_args(payload)?; to_json(fetch_usage(&args.account_id).await?) diff --git a/src/App.tsx b/src/App.tsx index 65ed2744..e80c2bb9 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -164,6 +164,7 @@ function App() { deleteAccount, renameAccount, importFromFile, + addFromAccessToken, exportAccountsSlimText, importAccountsSlimText, startOAuthLogin, @@ -1825,6 +1826,7 @@ function App() { isOpen={isAddModalOpen} onClose={() => setIsAddModalOpen(false)} onImportFile={importFromFile} + onAddAccessToken={addFromAccessToken} onStartOAuth={startOAuthLogin} onCompleteOAuth={completeOAuthLogin} onCancelOAuth={cancelOAuthLogin} diff --git a/src/components/AccountCard.tsx b/src/components/AccountCard.tsx index d06554db..dd30ae45 100644 --- a/src/components/AccountCard.tsx +++ b/src/components/AccountCard.tsx @@ -226,7 +226,9 @@ export function AccountCard({ ? account.plan_type.charAt(0).toUpperCase() + account.plan_type.slice(1) : account.auth_mode === "api_key" ? "API Key" - : "Unknown"; + : account.auth_mode === "codex_access_token" + ? "Access Token" + : "Unknown"; const planColors: Record = { pro: "bg-indigo-50 text-indigo-700 border-indigo-200 dark:bg-indigo-900/30 dark:text-indigo-300 dark:border-indigo-700", @@ -235,9 +237,10 @@ export function AccountCard({ enterprise: "bg-amber-50 text-amber-700 border-amber-200 dark:bg-amber-900/30 dark:text-amber-300 dark:border-amber-700", free: "bg-gray-50 text-gray-600 border-gray-200 dark:bg-gray-800 dark:text-gray-300 dark:border-gray-700", api_key: "bg-orange-50 text-orange-700 border-orange-200 dark:bg-orange-900/30 dark:text-orange-300 dark:border-orange-700", + codex_access_token: "bg-cyan-50 text-cyan-700 border-cyan-200 dark:bg-cyan-900/30 dark:text-cyan-300 dark:border-cyan-700", }; - const planKey = account.plan_type?.toLowerCase() || "api_key"; + const planKey = account.plan_type?.toLowerCase() || account.auth_mode; const planColorClass = planColors[planKey] || planColors.free; const showSubscriptionStatus = account.auth_mode === "chat_g_p_t"; const subscriptionStatus = getSubscriptionStatus(account.subscription_expires_at); diff --git a/src/components/AddAccountModal.tsx b/src/components/AddAccountModal.tsx index 9b795dec..49f251ae 100644 --- a/src/components/AddAccountModal.tsx +++ b/src/components/AddAccountModal.tsx @@ -11,17 +11,19 @@ interface AddAccountModalProps { isOpen: boolean; onClose: () => void; onImportFile: (source: FileSource, name: string) => Promise; + onAddAccessToken: (accessToken: string, name: string) => Promise; onStartOAuth: (name: string) => Promise<{ auth_url: string }>; onCompleteOAuth: () => Promise; onCancelOAuth: () => Promise; } -type Tab = "oauth" | "import"; +type Tab = "oauth" | "import" | "accessToken"; export function AddAccountModal({ isOpen, onClose, onImportFile, + onAddAccessToken, onStartOAuth, onCompleteOAuth, onCancelOAuth, @@ -29,6 +31,7 @@ export function AddAccountModal({ const [activeTab, setActiveTab] = useState("oauth"); const [name, setName] = useState(""); const [fileSource, setFileSource] = useState(null); + const [accessToken, setAccessToken] = useState(""); const [loading, setLoading] = useState(false); const [error, setError] = useState(null); const [oauthPending, setOauthPending] = useState(false); @@ -40,6 +43,7 @@ export function AddAccountModal({ const resetForm = () => { setName(""); setFileSource(null); + setAccessToken(""); setError(null); setLoading(false); setOauthPending(false); @@ -108,6 +112,49 @@ export function AddAccountModal({ } }; + const handleAddAccessToken = async () => { + if (!name.trim()) { + setError("Please enter an account name"); + return; + } + if (!accessToken.trim()) { + setError("Please paste a CODEX_ACCESS_TOKEN value"); + return; + } + + try { + setLoading(true); + setError(null); + await onAddAccessToken(accessToken.trim(), name.trim()); + handleClose(); + } catch (err) { + setError(err instanceof Error ? err.message : String(err)); + setLoading(false); + } + }; + + const handlePrimaryAction = () => { + if (activeTab === "oauth") { + void handleOAuthLogin(); + return; + } + + if (activeTab === "import") { + void handleImportFile(); + return; + } + + void handleAddAccessToken(); + }; + + const primaryLabel = loading + ? "Adding..." + : activeTab === "oauth" + ? "Generate Login Link" + : activeTab === "import" + ? "Import" + : "Add Token"; + if (!isOpen) return null; return ( @@ -126,11 +173,11 @@ export function AddAccountModal({ {/* Tabs */}
- {(["oauth", "import"] as Tab[]).map((tab) => ( + {(["oauth", "import", "accessToken"] as Tab[]).map((tab) => ( ))}
@@ -250,6 +301,23 @@ export function AddAccountModal({ )} + {activeTab === "accessToken" && ( +
+ + setAccessToken(e.target.value)} + placeholder="Paste token" + autoComplete="off" + spellCheck={false} + className="w-full px-4 py-2.5 bg-white dark:bg-gray-800 border border-gray-200 dark:border-gray-700 rounded-lg text-gray-900 dark:text-gray-100 placeholder-gray-400 dark:placeholder-gray-500 focus:outline-none focus:border-gray-400 dark:focus:border-gray-500 focus:ring-1 focus:ring-gray-400 dark:focus:ring-gray-500 transition-colors font-mono text-xs" + /> +
+ )} + {/* Error */} {error && (
@@ -267,15 +335,11 @@ export function AddAccountModal({ Cancel
diff --git a/src/hooks/useAccounts.ts b/src/hooks/useAccounts.ts index 94064a8f..6dd16149 100644 --- a/src/hooks/useAccounts.ts +++ b/src/hooks/useAccounts.ts @@ -288,6 +288,22 @@ export function useAccounts() { [loadAccounts, refreshUsage] ); + const addFromAccessToken = useCallback( + async (accessToken: string, name: string) => { + try { + await invokeBackend("add_account_from_access_token", { + name, + accessToken, + }); + const accountList = await loadAccounts(); + await refreshUsage(accountList); + } catch (err) { + throw err; + } + }, + [loadAccounts, refreshUsage] + ); + const startOAuthLogin = useCallback(async (accountName: string) => { try { const info = await invokeBackend<{ auth_url: string; callback_port: number }>( @@ -426,6 +442,7 @@ export function useAccounts() { deleteAccount, renameAccount, importFromFile, + addFromAccessToken, exportAccountsSlimText, importAccountsSlimText, exportAccountsFullEncryptedFile, diff --git a/src/types/index.ts b/src/types/index.ts index 57c30bd8..986c764b 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -1,6 +1,6 @@ // Types matching the Rust backend -export type AuthMode = "api_key" | "chat_g_p_t"; +export type AuthMode = "api_key" | "chat_g_p_t" | "codex_access_token"; export type DockDisplayMode = "show_in_dock" | "menu_bar_only"; export interface AccountInfo { From 223fa2f455be00e37d58c660e584c009a634a13a Mon Sep 17 00:00:00 2001 From: zakyirsyaad Date: Sun, 5 Jul 2026 20:31:24 +0700 Subject: [PATCH 2/8] fix: clarify K12 access token modal --- src/components/AddAccountModal.tsx | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/src/components/AddAccountModal.tsx b/src/components/AddAccountModal.tsx index 49f251ae..f4b1acc6 100644 --- a/src/components/AddAccountModal.tsx +++ b/src/components/AddAccountModal.tsx @@ -118,7 +118,7 @@ export function AddAccountModal({ return; } if (!accessToken.trim()) { - setError("Please paste a CODEX_ACCESS_TOKEN value"); + setError("Please paste a K12 CODEX_ACCESS_TOKEN value"); return; } @@ -153,7 +153,7 @@ export function AddAccountModal({ ? "Generate Login Link" : activeTab === "import" ? "Import" - : "Add Token"; + : "Add K12 Token"; if (!isOpen) return null; @@ -196,7 +196,7 @@ export function AddAccountModal({ ? "ChatGPT Login" : tab === "import" ? "Import File" - : "Access Token"} + : "K12 Token"} ))} @@ -302,15 +302,20 @@ export function AddAccountModal({ )} {activeTab === "accessToken" && ( -
+
+
+ This option is only for K12 accounts that provide a Codex + CODEX_ACCESS_TOKEN. + Use ChatGPT Login or Import File for other account types. +
setAccessToken(e.target.value)} - placeholder="Paste token" + placeholder="Paste K12 access token" autoComplete="off" spellCheck={false} className="w-full px-4 py-2.5 bg-white dark:bg-gray-800 border border-gray-200 dark:border-gray-700 rounded-lg text-gray-900 dark:text-gray-100 placeholder-gray-400 dark:placeholder-gray-500 focus:outline-none focus:border-gray-400 dark:focus:border-gray-500 focus:ring-1 focus:ring-gray-400 dark:focus:ring-gray-500 transition-colors font-mono text-xs" From e136b0c8a7c87fe987aa5fa59750931c3981e407 Mon Sep 17 00:00:00 2001 From: zakyirsyaad Date: Tue, 7 Jul 2026 01:13:22 +0700 Subject: [PATCH 3/8] fix: resolve Codex CLI path for token login --- src-tauri/src/auth/switcher.rs | 101 ++++++++++++++++++++++++++++++++- src/App.tsx | 2 + 2 files changed, 100 insertions(+), 3 deletions(-) diff --git a/src-tauri/src/auth/switcher.rs b/src-tauri/src/auth/switcher.rs index 94f4b3d0..dfe9c980 100644 --- a/src-tauri/src/auth/switcher.rs +++ b/src-tauri/src/auth/switcher.rs @@ -2,7 +2,7 @@ use std::fs; use std::io::Write; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::process::{Command, Stdio}; use anyhow::{Context, Result}; @@ -122,7 +122,7 @@ fn login_with_codex_access_token(token: &str) -> Result<()> { anyhow::bail!("Codex access token is empty"); } - let mut child = Command::new("codex") + let mut child = Command::new(resolve_codex_executable()) .args(["login", "--with-access-token"]) .stdin(Stdio::piped()) .stdout(Stdio::piped()) @@ -159,6 +159,53 @@ fn login_with_codex_access_token(token: &str) -> Result<()> { Ok(()) } +/// Resolve the `codex` CLI executable, falling back to well-known install +/// locations when it isn't found on the current process's PATH. +/// +/// GUI apps launched from Finder/Dock/Spotlight on macOS inherit launchd's +/// minimal PATH, not the user's shell PATH — so a Homebrew-installed `codex` +/// (e.g. under `/opt/homebrew/bin`) is invisible to `Command::new("codex")` +/// even though it works fine from a terminal. +fn resolve_codex_executable() -> PathBuf { + if let Some(path) = find_on_path("codex", std::env::var_os("PATH")) { + return path; + } + + #[cfg(unix)] + { + for candidate in fallback_codex_paths(dirs::home_dir().as_deref()) { + if candidate.is_file() { + return candidate; + } + } + } + + PathBuf::from("codex") +} + +fn find_on_path(name: &str, path_var: Option) -> Option { + let path_var = path_var?; + std::env::split_paths(&path_var).find_map(|dir| { + let candidate = dir.join(name); + candidate.is_file().then_some(candidate) + }) +} + +#[cfg(unix)] +fn fallback_codex_paths(home: Option<&Path>) -> Vec { + let mut candidates = vec![ + PathBuf::from("/opt/homebrew/bin/codex"), + PathBuf::from("/usr/local/bin/codex"), + ]; + + if let Some(home) = home { + candidates.push(home.join(".local/bin/codex")); + candidates.push(home.join(".npm-global/bin/codex")); + } + + candidates +} + fn redact_access_token_from_output(output: &str, token: &str) -> String { if token.is_empty() { return output.to_string(); @@ -169,7 +216,55 @@ fn redact_access_token_from_output(output: &str, token: &str) -> String { #[cfg(test)] mod tests { - use super::{create_access_token_auth_json, redact_access_token_from_output}; + use super::{create_access_token_auth_json, find_on_path, redact_access_token_from_output}; + + #[test] + fn find_on_path_returns_none_when_path_is_unset() { + assert_eq!(find_on_path("codex", None), None); + } + + #[test] + fn find_on_path_returns_none_when_not_present_in_any_directory() { + let path_var = std::ffi::OsString::from("/definitely/does/not/exist/xyz"); + assert_eq!(find_on_path("codex", Some(path_var)), None); + } + + #[test] + fn find_on_path_locates_executable_within_a_path_directory() { + let dir = std::env::temp_dir().join(format!( + "codex-switcher-find-on-path-test-{}", + std::process::id() + )); + std::fs::create_dir_all(&dir).unwrap(); + let exe = dir.join("codex"); + std::fs::write(&exe, b"").unwrap(); + + let path_var = std::ffi::OsString::from(dir.as_os_str()); + let found = find_on_path("codex", Some(path_var)); + + std::fs::remove_dir_all(&dir).ok(); + + assert_eq!(found.as_deref(), Some(exe.as_path())); + } + + #[cfg(unix)] + #[test] + fn fallback_codex_paths_includes_common_homebrew_prefixes() { + let candidates = super::fallback_codex_paths(None); + + assert!(candidates.contains(&std::path::PathBuf::from("/opt/homebrew/bin/codex"))); + assert!(candidates.contains(&std::path::PathBuf::from("/usr/local/bin/codex"))); + } + + #[cfg(unix)] + #[test] + fn fallback_codex_paths_includes_home_relative_locations_when_home_is_known() { + let home = std::path::PathBuf::from("/Users/example"); + let candidates = super::fallback_codex_paths(Some(&home)); + + assert!(candidates.contains(&home.join(".local/bin/codex"))); + assert!(candidates.contains(&home.join(".npm-global/bin/codex"))); + } #[test] fn redacts_access_token_from_cli_error_output() { diff --git a/src/App.tsx b/src/App.tsx index e80c2bb9..46f562c9 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -502,6 +502,8 @@ function App() { await switchAccount(accountId); } catch (err) { console.error("Failed to switch account:", err); + const accountName = accounts.find((account) => account.id === accountId)?.name ?? "account"; + showWarmupToast(`Switch failed for ${accountName}: ${formatWarmupError(err)}`, true); } finally { setSwitchingId(null); } From ed621cd26630a0c3cf616fdc697001ed0dcc0870 Mon Sep 17 00:00:00 2001 From: zakyirsyaad Date: Tue, 7 Jul 2026 12:12:38 +0700 Subject: [PATCH 4/8] fix: write auth files atomically --- src-tauri/src/auth/storage.rs | 114 +++++++++++++++++++++++++++++---- src-tauri/src/auth/switcher.rs | 11 +--- 2 files changed, 104 insertions(+), 21 deletions(-) diff --git a/src-tauri/src/auth/storage.rs b/src-tauri/src/auth/storage.rs index 54514851..a92074ec 100644 --- a/src-tauri/src/auth/storage.rs +++ b/src-tauri/src/auth/storage.rs @@ -5,6 +5,7 @@ use std::path::PathBuf; use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; +use uuid::Uuid; use crate::types::{AccountsStore, AppSettings, AuthData, StoredAccount}; @@ -65,16 +66,55 @@ pub fn save_app_settings(settings: &AppSettings) -> Result<()> { } let content = serde_json::to_string_pretty(settings).context("Failed to serialize settings")?; - fs::write(&path, content) + write_file_atomically(&path, &content) .with_context(|| format!("Failed to write settings file: {}", path.display()))?; + Ok(()) +} + +/// Write `content` to `path` by writing a sibling temp file and atomically +/// renaming it into place, instead of writing directly to `path`. +/// +/// `fs::write` truncates the target file before writing the new bytes, so any +/// concurrent reader (the tray's 1s file watcher, the 60s usage poller, a +/// separate `codex-web` LAN process, or the official `codex` CLI reading +/// `auth.json`) that opens the file during that window observes a truncated or +/// empty file. `fs::rename` within the same directory is atomic, so readers +/// always see either the fully-old or fully-new content, never a partial +/// write. The temp filename is suffixed with a UUID rather than a PID so two +/// writers on the same process (e.g. a tray-triggered switch racing a +/// main-window switch) never collide on the same temp file. Shared across +/// `auth::storage` and `auth::switcher`, which both persist small JSON/text +/// files that outside readers can observe mid-write. +pub(crate) fn write_file_atomically(path: &std::path::Path, content: &str) -> Result<()> { + let extension = path.extension().and_then(|e| e.to_str()).unwrap_or("tmp"); + let tmp_path = path.with_extension(format!("{extension}.tmp.{}", Uuid::new_v4())); + + let result = write_temp_then_rename(&tmp_path, path, content); + if result.is_err() { + let _ = fs::remove_file(&tmp_path); + } + result +} + +fn write_temp_then_rename( + tmp_path: &std::path::Path, + path: &std::path::Path, + content: &str, +) -> Result<()> { + fs::write(tmp_path, content) + .with_context(|| format!("Failed to write temp file: {}", tmp_path.display()))?; + #[cfg(unix)] { use std::os::unix::fs::PermissionsExt; - let perms = fs::Permissions::from_mode(0o600); - fs::set_permissions(&path, perms)?; + fs::set_permissions(tmp_path, fs::Permissions::from_mode(0o600)) + .with_context(|| format!("Failed to set permissions: {}", tmp_path.display()))?; } + fs::rename(tmp_path, path) + .with_context(|| format!("Failed to finalize file: {}", path.display()))?; + Ok(()) } @@ -91,17 +131,9 @@ pub fn save_accounts(store: &AccountsStore) -> Result<()> { let content = serde_json::to_string_pretty(store).context("Failed to serialize accounts store")?; - fs::write(&path, content) + write_file_atomically(&path, &content) .with_context(|| format!("Failed to write accounts file: {}", path.display()))?; - // Set restrictive permissions on Unix - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let perms = fs::Permissions::from_mode(0o600); - fs::set_permissions(&path, perms)?; - } - Ok(()) } @@ -305,3 +337,61 @@ pub fn set_masked_account_ids(ids: Vec) -> Result<()> { save_accounts(&store)?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::Arc; + + /// Reproduces the "Failed to parse accounts file" bug: a reader (e.g. the + /// tray's 1s file watcher) can open the file while a writer is mid-`fs::write`, + /// since plain `fs::write` truncates before the new bytes land. + #[test] + fn concurrent_readers_never_observe_a_torn_write() { + let dir = std::env::temp_dir().join(format!("codex-switcher-test-{}", Uuid::new_v4())); + fs::create_dir_all(&dir).unwrap(); + let path = dir.join("accounts.json"); + + let payload_a = format!("{{\"marker\":\"A\",\"pad\":\"{}\"}}", "a".repeat(200_000)); + let payload_b = format!("{{\"marker\":\"B\",\"pad\":\"{}\"}}", "b".repeat(200_000)); + + write_file_atomically(&path, &payload_a).unwrap(); + + let stop = Arc::new(AtomicBool::new(false)); + let writer = { + let stop = stop.clone(); + let path = path.clone(); + let payload_a = payload_a.clone(); + let payload_b = payload_b.clone(); + std::thread::spawn(move || { + for i in 0..200 { + let payload = if i % 2 == 0 { &payload_b } else { &payload_a }; + write_file_atomically(&path, payload).unwrap(); + } + stop.store(true, Ordering::SeqCst); + }) + }; + + let mut reads = 0; + while !stop.load(Ordering::SeqCst) { + if let Ok(content) = fs::read_to_string(&path) { + assert!( + content == payload_a || content == payload_b, + "torn read detected: got {} bytes (expected {} or {})", + content.len(), + payload_a.len(), + payload_b.len() + ); + reads += 1; + } + } + writer.join().unwrap(); + assert!( + reads > 0, + "test didn't overlap reads with writes — increase iterations" + ); + + let _ = fs::remove_dir_all(&dir); + } +} diff --git a/src-tauri/src/auth/switcher.rs b/src-tauri/src/auth/switcher.rs index dfe9c980..f92583ce 100644 --- a/src-tauri/src/auth/switcher.rs +++ b/src-tauri/src/auth/switcher.rs @@ -8,6 +8,7 @@ use std::process::{Command, Stdio}; use anyhow::{Context, Result}; use chrono::Utc; +use crate::auth::storage::write_file_atomically; use crate::types::{ parse_chatgpt_id_token_claims, AuthData, AuthDotJson, StoredAccount, TokenData, }; @@ -46,17 +47,9 @@ pub fn switch_to_account(account: &StoredAccount) -> Result<()> { let content = serde_json::to_string_pretty(&auth_json).context("Failed to serialize auth.json")?; - fs::write(&auth_path, content) + write_file_atomically(&auth_path, &content) .with_context(|| format!("Failed to write auth.json: {}", auth_path.display()))?; - // Set restrictive permissions on Unix - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let perms = fs::Permissions::from_mode(0o600); - fs::set_permissions(&auth_path, perms)?; - } - Ok(()) } From dfd6701101b5851594427eb5e90a3b08993266b6 Mon Sep 17 00:00:00 2001 From: zakyirsyaad Date: Thu, 9 Jul 2026 01:23:44 +0700 Subject: [PATCH 5/8] fix: serialize account store mutations --- src-tauri/src/auth/storage.rs | 541 ++++++++++++++++++---------- src-tauri/src/auth/token_refresh.rs | 67 +++- src-tauri/src/commands/account.rs | 63 ++-- src-tauri/src/commands/oauth.rs | 19 +- 4 files changed, 442 insertions(+), 248 deletions(-) diff --git a/src-tauri/src/auth/storage.rs b/src-tauri/src/auth/storage.rs index a92074ec..6ec1f356 100644 --- a/src-tauri/src/auth/storage.rs +++ b/src-tauri/src/auth/storage.rs @@ -1,13 +1,25 @@ //! Account storage module - manages reading and writing accounts.json - -use std::fs; -use std::path::PathBuf; +//! +//! Concurrency model: many concurrent readers exist (the tray's 1s file +//! watcher, the 60s usage poller, Tauri commands, and a separate `codex-web` +//! process), so every write must be atomic (`write_file_atomically`), and +//! every read-modify-write cycle must run under the exclusive store lock +//! (`mutate_accounts`) or concurrent writers silently overwrite each other's +//! changes. + +use std::fs::{self, File, OpenOptions}; +use std::io::Write; +use std::path::{Path, PathBuf}; use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; use uuid::Uuid; -use crate::types::{AccountsStore, AppSettings, AuthData, StoredAccount}; +use crate::types::{AccountsStore, AppSettings, StoredAccount}; + +const ACCOUNTS_FILE: &str = "accounts.json"; +const SETTINGS_FILE: &str = "settings.json"; +const STORE_LOCK_FILE: &str = ".store.lock"; /// Get the path to the codex-switcher config directory pub fn get_config_dir() -> Result { @@ -17,16 +29,45 @@ pub fn get_config_dir() -> Result { /// Get the path to accounts.json pub fn get_accounts_file() -> Result { - Ok(get_config_dir()?.join("accounts.json")) + Ok(get_config_dir()?.join(ACCOUNTS_FILE)) } pub fn get_settings_file() -> Result { - Ok(get_config_dir()?.join("settings.json")) + Ok(get_config_dir()?.join(SETTINGS_FILE)) } -/// Load the accounts store from disk +/// Acquire the exclusive advisory lock that serializes every store write, +/// both across threads in this process and across processes (the desktop app +/// and the `codex-web` LAN server share the same files). Blocks until the +/// lock is available; the OS releases it automatically when the returned +/// `File` drops, including when the holding process crashes. +fn acquire_store_lock(dir: &Path) -> Result { + fs::create_dir_all(dir) + .with_context(|| format!("Failed to create config directory: {}", dir.display()))?; + + let lock_path = dir.join(STORE_LOCK_FILE); + let lock_file = OpenOptions::new() + .create(true) + .read(true) + .append(true) + .open(&lock_path) + .with_context(|| format!("Failed to open store lock file: {}", lock_path.display()))?; + + lock_file + .lock() + .with_context(|| format!("Failed to lock store: {}", lock_path.display()))?; + + Ok(lock_file) +} + +/// Load the accounts store from disk. Lock-free: atomic writes guarantee a +/// reader always sees a complete file. pub fn load_accounts() -> Result { - let path = get_accounts_file()?; + load_accounts_in(&get_config_dir()?) +} + +fn load_accounts_in(dir: &Path) -> Result { + let path = dir.join(ACCOUNTS_FILE); if !path.exists() { return Ok(AccountsStore::default()); @@ -35,6 +76,13 @@ pub fn load_accounts() -> Result { let content = fs::read_to_string(&path) .with_context(|| format!("Failed to read accounts file: {}", path.display()))?; + // An empty file carries no data worth preserving; treat it like a missing + // file so a past crash can't brick every launch, but keep failing loudly + // on non-empty garbage — overwriting that would destroy real accounts. + if content.trim().is_empty() { + return Ok(AccountsStore::default()); + } + let store: AccountsStore = serde_json::from_str(&content) .with_context(|| format!("Failed to parse accounts file: {}", path.display()))?; @@ -51,42 +99,75 @@ pub fn load_app_settings() -> Result { let content = fs::read_to_string(&path) .with_context(|| format!("Failed to read settings file: {}", path.display()))?; + if content.trim().is_empty() { + return Ok(AppSettings::default()); + } + let settings: AppSettings = serde_json::from_str(&content) .with_context(|| format!("Failed to parse settings file: {}", path.display()))?; Ok(settings) } +/// Persist app settings. Settings load-modify-save cycles all run on the main +/// thread (menu and Dock handlers), so only the write itself takes the store +/// lock — enough to stay safe against writers in other processes. pub fn save_app_settings(settings: &AppSettings) -> Result<()> { - let path = get_settings_file()?; - - if let Some(parent) = path.parent() { - fs::create_dir_all(parent) - .with_context(|| format!("Failed to create config directory: {}", parent.display()))?; - } + let dir = get_config_dir()?; + let _lock = acquire_store_lock(&dir)?; let content = serde_json::to_string_pretty(settings).context("Failed to serialize settings")?; - write_file_atomically(&path, &content) - .with_context(|| format!("Failed to write settings file: {}", path.display()))?; + write_file_atomically(&dir.join(SETTINGS_FILE), &content) + .with_context(|| format!("Failed to write settings file: {}", dir.display())) +} - Ok(()) +/// Run a read-modify-write cycle on the accounts store under the exclusive +/// store lock: load, apply `mutate`, save. If `mutate` returns an error, +/// nothing is written. +/// +/// The closure must not call any other locking storage function +/// (`mutate_accounts`, `save_app_settings`, `add_account`, ...) — the lock is +/// not re-entrant and doing so deadlocks. Lock-free helpers like +/// `switch_to_account` and plain access to `store` are fine. +pub fn mutate_accounts(mutate: impl FnOnce(&mut AccountsStore) -> Result) -> Result { + mutate_accounts_in(&get_config_dir()?, mutate) +} + +fn mutate_accounts_in( + dir: &Path, + mutate: impl FnOnce(&mut AccountsStore) -> Result, +) -> Result { + let _lock = acquire_store_lock(dir)?; + + let mut store = load_accounts_in(dir)?; + let value = mutate(&mut store)?; + + let content = + serde_json::to_string_pretty(&store).context("Failed to serialize accounts store")?; + write_file_atomically(&dir.join(ACCOUNTS_FILE), &content) + .with_context(|| format!("Failed to write accounts file: {}", dir.display()))?; + + Ok(value) } /// Write `content` to `path` by writing a sibling temp file and atomically /// renaming it into place, instead of writing directly to `path`. /// -/// `fs::write` truncates the target file before writing the new bytes, so any +/// `fs::write` truncates the target before writing the new bytes, so any /// concurrent reader (the tray's 1s file watcher, the 60s usage poller, a /// separate `codex-web` LAN process, or the official `codex` CLI reading -/// `auth.json`) that opens the file during that window observes a truncated or -/// empty file. `fs::rename` within the same directory is atomic, so readers -/// always see either the fully-old or fully-new content, never a partial -/// write. The temp filename is suffixed with a UUID rather than a PID so two -/// writers on the same process (e.g. a tray-triggered switch racing a -/// main-window switch) never collide on the same temp file. Shared across -/// `auth::storage` and `auth::switcher`, which both persist small JSON/text -/// files that outside readers can observe mid-write. -pub(crate) fn write_file_atomically(path: &std::path::Path, content: &str) -> Result<()> { +/// `auth.json`) that opens the file during that window observes a truncated +/// or empty file. A same-directory `rename` is atomic, so readers always see +/// either the fully-old or fully-new content. The temp file is created with +/// 0600 from the first byte (these files carry credentials) and fsynced +/// before the rename so a crash can't leave the target pointing at data that +/// never reached disk. +/// +/// Callers are expected to hold the store lock; the stale-temp sweep assumes +/// no other writer of ours is mid-flight on the same target. +pub(crate) fn write_file_atomically(path: &Path, content: &str) -> Result<()> { + sweep_stale_tmp_files(path); + let extension = path.extension().and_then(|e| e.to_str()).unwrap_or("tmp"); let tmp_path = path.with_extension(format!("{extension}.tmp.{}", Uuid::new_v4())); @@ -97,99 +178,113 @@ pub(crate) fn write_file_atomically(path: &std::path::Path, content: &str) -> Re result } -fn write_temp_then_rename( - tmp_path: &std::path::Path, - path: &std::path::Path, - content: &str, -) -> Result<()> { - fs::write(tmp_path, content) - .with_context(|| format!("Failed to write temp file: {}", tmp_path.display()))?; +/// Best-effort removal of `.tmp.` leftovers from writes that +/// crashed between creating the temp file and renaming it into place. +fn sweep_stale_tmp_files(path: &Path) { + let Some(parent) = path.parent() else { return }; + let Some(name) = path.file_name().and_then(|n| n.to_str()) else { + return; + }; + let prefix = format!("{name}.tmp."); + let Ok(entries) = fs::read_dir(parent) else { + return; + }; + for entry in entries.flatten() { + if entry + .file_name() + .to_str() + .is_some_and(|n| n.starts_with(&prefix)) + { + let _ = fs::remove_file(entry.path()); + } + } +} + +fn write_temp_then_rename(tmp_path: &Path, path: &Path, content: &str) -> Result<()> { + let mut options = OpenOptions::new(); + options.write(true).create_new(true); #[cfg(unix)] { - use std::os::unix::fs::PermissionsExt; - fs::set_permissions(tmp_path, fs::Permissions::from_mode(0o600)) - .with_context(|| format!("Failed to set permissions: {}", tmp_path.display()))?; + use std::os::unix::fs::OpenOptionsExt; + options.mode(0o600); } - fs::rename(tmp_path, path) + let mut file = options + .open(tmp_path) + .with_context(|| format!("Failed to create temp file: {}", tmp_path.display()))?; + file.write_all(content.as_bytes()) + .with_context(|| format!("Failed to write temp file: {}", tmp_path.display()))?; + file.sync_all() + .with_context(|| format!("Failed to flush temp file: {}", tmp_path.display()))?; + drop(file); + + rename_replacing(tmp_path, path) .with_context(|| format!("Failed to finalize file: {}", path.display()))?; Ok(()) } -/// Save the accounts store to disk -pub fn save_accounts(store: &AccountsStore) -> Result<()> { - let path = get_accounts_file()?; +#[cfg(not(windows))] +fn rename_replacing(from: &Path, to: &Path) -> std::io::Result<()> { + fs::rename(from, to) +} - // Ensure the config directory exists - if let Some(parent) = path.parent() { - fs::create_dir_all(parent) - .with_context(|| format!("Failed to create config directory: {}", parent.display()))?; +/// On Windows, replacing a file that another process holds open without +/// FILE_SHARE_DELETE (antivirus scanners, non-Rust readers) fails with a +/// sharing violation that clears once the reader closes — retry briefly. +#[cfg(windows)] +fn rename_replacing(from: &Path, to: &Path) -> std::io::Result<()> { + const ATTEMPTS: u32 = 5; + let mut last_error = None; + for attempt in 0..ATTEMPTS { + if attempt > 0 { + std::thread::sleep(std::time::Duration::from_millis(20)); + } + match fs::rename(from, to) { + Ok(()) => return Ok(()), + Err(error) => last_error = Some(error), + } } - - let content = - serde_json::to_string_pretty(store).context("Failed to serialize accounts store")?; - - write_file_atomically(&path, &content) - .with_context(|| format!("Failed to write accounts file: {}", path.display()))?; - - Ok(()) + Err(last_error.expect("rename attempted at least once")) } /// Add a new account to the store pub fn add_account(account: StoredAccount) -> Result { - let mut store = load_accounts()?; - - // Check for duplicate names - if store.accounts.iter().any(|a| a.name == account.name) { - anyhow::bail!("An account with name '{}' already exists", account.name); - } + mutate_accounts(|store| { + if store.accounts.iter().any(|a| a.name == account.name) { + anyhow::bail!("An account with name '{}' already exists", account.name); + } - let account_clone = account.clone(); - store.accounts.push(account); + let account_clone = account.clone(); + store.accounts.push(account); - // If this is the first account, make it active - if store.accounts.len() == 1 { - store.active_account_id = Some(account_clone.id.clone()); - } + // If this is the first account, make it active + if store.accounts.len() == 1 { + store.active_account_id = Some(account_clone.id.clone()); + } - save_accounts(&store)?; - Ok(account_clone) + Ok(account_clone) + }) } /// Remove an account by ID pub fn remove_account(account_id: &str) -> Result<()> { - let mut store = load_accounts()?; - - let initial_len = store.accounts.len(); - store.accounts.retain(|a| a.id != account_id); - - if store.accounts.len() == initial_len { - anyhow::bail!("Account not found: {account_id}"); - } - - // If we removed the active account, clear it or set to first available - if store.active_account_id.as_deref() == Some(account_id) { - store.active_account_id = store.accounts.first().map(|a| a.id.clone()); - } + mutate_accounts(|store| { + let initial_len = store.accounts.len(); + store.accounts.retain(|a| a.id != account_id); - save_accounts(&store)?; - Ok(()) -} - -/// Update the active account ID -pub fn set_active_account(account_id: &str) -> Result<()> { - let mut store = load_accounts()?; + if store.accounts.len() == initial_len { + anyhow::bail!("Account not found: {account_id}"); + } - // Verify the account exists - if !store.accounts.iter().any(|a| a.id == account_id) { - anyhow::bail!("Account not found: {account_id}"); - } + // If we removed the active account, clear it or set to first available + if store.active_account_id.as_deref() == Some(account_id) { + store.active_account_id = store.accounts.first().map(|a| a.id.clone()); + } - store.active_account_id = Some(account_id.to_string()); - save_accounts(&store)?; - Ok(()) + Ok(()) + }) } /// Get an account by ID @@ -208,18 +303,6 @@ pub fn get_active_account() -> Result> { Ok(store.accounts.into_iter().find(|a| a.id == *active_id)) } -/// Update an account's last_used_at timestamp -pub fn touch_account(account_id: &str) -> Result<()> { - let mut store = load_accounts()?; - - if let Some(account) = store.accounts.iter_mut().find(|a| a.id == account_id) { - account.last_used_at = Some(chrono::Utc::now()); - save_accounts(&store)?; - } - - Ok(()) -} - /// Update an account's metadata (name, email, plan_type, subscription expiry) pub fn update_account_metadata( account_id: &str, @@ -228,100 +311,43 @@ pub fn update_account_metadata( plan_type: Option, subscription_expires_at: Option>>, ) -> Result { - let mut store = load_accounts()?; - - // Check for duplicate names first (if renaming) - if let Some(ref new_name) = name { - if store - .accounts - .iter() - .any(|a| a.id != account_id && a.name == *new_name) - { - anyhow::bail!("An account with name '{new_name}' already exists"); + mutate_accounts(|store| { + // Check for duplicate names first (if renaming) + if let Some(ref new_name) = name { + if store + .accounts + .iter() + .any(|a| a.id != account_id && a.name == *new_name) + { + anyhow::bail!("An account with name '{new_name}' already exists"); + } } - } - - // Now find and update the account - let account = store - .accounts - .iter_mut() - .find(|a| a.id == account_id) - .context("Account not found")?; - - if let Some(new_name) = name { - account.name = new_name; - } - - if email.is_some() { - account.email = email; - } - - if plan_type.is_some() { - account.plan_type = plan_type; - } - - if let Some(subscription_expires_at) = subscription_expires_at { - account.subscription_expires_at = subscription_expires_at; - } - let updated = account.clone(); - save_accounts(&store)?; - Ok(updated) -} + // Now find and update the account + let account = store + .accounts + .iter_mut() + .find(|a| a.id == account_id) + .context("Account not found")?; -/// Update ChatGPT OAuth tokens for an account and return the updated account. -pub fn update_account_chatgpt_tokens( - account_id: &str, - id_token: String, - access_token: String, - refresh_token: String, - chatgpt_account_id: Option, - email: Option, - plan_type: Option, - subscription_expires_at: Option>, -) -> Result { - let mut store = load_accounts()?; - - let account = store - .accounts - .iter_mut() - .find(|a| a.id == account_id) - .context("Account not found")?; - - match &mut account.auth_data { - AuthData::ChatGPT { - id_token: stored_id_token, - access_token: stored_access_token, - refresh_token: stored_refresh_token, - account_id: stored_account_id, - } => { - *stored_id_token = id_token; - *stored_access_token = access_token; - *stored_refresh_token = refresh_token; - if let Some(new_account_id) = chatgpt_account_id { - *stored_account_id = Some(new_account_id); - } - } - AuthData::ApiKey { .. } | AuthData::CodexAccessToken { .. } => { - anyhow::bail!("Cannot update OAuth tokens for this account type"); + if let Some(new_name) = name { + account.name = new_name; } - } - if let Some(new_email) = email { - account.email = Some(new_email); - } + if email.is_some() { + account.email = email; + } - if let Some(new_plan_type) = plan_type { - account.plan_type = Some(new_plan_type); - } + if plan_type.is_some() { + account.plan_type = plan_type; + } - if let Some(subscription_expires_at) = subscription_expires_at { - account.subscription_expires_at = Some(subscription_expires_at); - } + if let Some(subscription_expires_at) = subscription_expires_at { + account.subscription_expires_at = subscription_expires_at; + } - let updated = account.clone(); - save_accounts(&store)?; - Ok(updated) + Ok(account.clone()) + }) } /// Get the list of masked account IDs @@ -332,10 +358,10 @@ pub fn get_masked_account_ids() -> Result> { /// Set the list of masked account IDs pub fn set_masked_account_ids(ids: Vec) -> Result<()> { - let mut store = load_accounts()?; - store.masked_account_ids = ids; - save_accounts(&store)?; - Ok(()) + mutate_accounts(|store| { + store.masked_account_ids = ids; + Ok(()) + }) } #[cfg(test)] @@ -344,13 +370,18 @@ mod tests { use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; + fn temp_store_dir() -> PathBuf { + let dir = std::env::temp_dir().join(format!("codex-switcher-test-{}", Uuid::new_v4())); + fs::create_dir_all(&dir).unwrap(); + dir + } + /// Reproduces the "Failed to parse accounts file" bug: a reader (e.g. the /// tray's 1s file watcher) can open the file while a writer is mid-`fs::write`, /// since plain `fs::write` truncates before the new bytes land. #[test] fn concurrent_readers_never_observe_a_torn_write() { - let dir = std::env::temp_dir().join(format!("codex-switcher-test-{}", Uuid::new_v4())); - fs::create_dir_all(&dir).unwrap(); + let dir = temp_store_dir(); let path = dir.join("accounts.json"); let payload_a = format!("{{\"marker\":\"A\",\"pad\":\"{}\"}}", "a".repeat(200_000)); @@ -394,4 +425,118 @@ mod tests { let _ = fs::remove_dir_all(&dir); } + + /// Reproduces the lost-update bug: unlocked load-modify-save cycles (e.g. + /// the 10-way-concurrent metadata updates in `refresh_all_accounts_usage`) + /// silently drop each other's changes. + #[test] + fn concurrent_mutations_never_lose_updates() { + let dir = temp_store_dir(); + const THREADS: usize = 8; + const ITERATIONS: usize = 25; + + let handles: Vec<_> = (0..THREADS) + .map(|thread| { + let dir = dir.clone(); + std::thread::spawn(move || { + for i in 0..ITERATIONS { + mutate_accounts_in(&dir, |store| { + let account = StoredAccount::new_api_key( + format!("account-{thread}-{i}"), + "dummy-api-key".to_string(), + ); + store.accounts.push(account); + Ok(()) + }) + .unwrap(); + } + }) + }) + .collect(); + + for handle in handles { + handle.join().unwrap(); + } + + let store = load_accounts_in(&dir).unwrap(); + assert_eq!( + store.accounts.len(), + THREADS * ITERATIONS, + "lost updates: concurrent mutations overwrote each other" + ); + + let _ = fs::remove_dir_all(&dir); + } + + #[test] + fn treats_empty_accounts_file_as_missing() { + let dir = temp_store_dir(); + fs::write(dir.join("accounts.json"), "").unwrap(); + + let store = load_accounts_in(&dir).unwrap(); + assert!(store.accounts.is_empty()); + assert_eq!(store.version, 1); + + let _ = fs::remove_dir_all(&dir); + } + + #[test] + fn rejects_corrupt_accounts_file_instead_of_wiping_it() { + let dir = temp_store_dir(); + fs::write(dir.join("accounts.json"), "{\"version\":1,").unwrap(); + + assert!(load_accounts_in(&dir).is_err()); + + let _ = fs::remove_dir_all(&dir); + } + + #[test] + fn failed_mutation_leaves_store_untouched() { + let dir = temp_store_dir(); + mutate_accounts_in(&dir, |store| { + store.accounts.push(StoredAccount::new_api_key( + "keep".into(), + "dummy-api-key".into(), + )); + Ok(()) + }) + .unwrap(); + + let result: Result<()> = mutate_accounts_in(&dir, |store| { + store.accounts.clear(); + anyhow::bail!("boom") + }); + assert!(result.is_err()); + + let store = load_accounts_in(&dir).unwrap(); + assert_eq!(store.accounts.len(), 1, "failed mutation must not be saved"); + + let _ = fs::remove_dir_all(&dir); + } + + #[test] + fn sweeps_stale_tmp_files_and_leaves_none_behind() { + let dir = temp_store_dir(); + let path = dir.join("accounts.json"); + + fs::write(dir.join("accounts.json.tmp.stale-crash-leftover"), "junk").unwrap(); + fs::write(dir.join("unrelated.txt"), "keep me").unwrap(); + + write_file_atomically(&path, "{\"ok\":true}").unwrap(); + + let names: Vec = fs::read_dir(&dir) + .unwrap() + .flatten() + .map(|e| e.file_name().to_string_lossy().into_owned()) + .collect(); + + assert!(names.contains(&"accounts.json".to_string())); + assert!(names.contains(&"unrelated.txt".to_string())); + assert!( + !names.iter().any(|n| n.contains(".tmp.")), + "stale or leftover temp files remain: {names:?}" + ); + + let _ = fs::remove_dir_all(&dir); + } } diff --git a/src-tauri/src/auth/token_refresh.rs b/src-tauri/src/auth/token_refresh.rs index c6b91842..593b42f4 100644 --- a/src-tauri/src/auth/token_refresh.rs +++ b/src-tauri/src/auth/token_refresh.rs @@ -5,7 +5,7 @@ use base64::Engine; use chrono::Utc; use tokio::time::{sleep, Duration}; -use super::{load_accounts, switch_to_account, update_account_chatgpt_tokens}; +use super::{mutate_accounts, switch_to_account}; use crate::types::{parse_chatgpt_id_token_claims, AuthData, StoredAccount}; const DEFAULT_ISSUER: &str = "https://auth.openai.com"; @@ -65,27 +65,58 @@ pub async fn refresh_chatgpt_tokens(account: &StoredAccount) -> Result { + *id_token = next_id_token; + *access_token = refreshed.access_token; + *refresh_token = next_refresh_token; + if let Some(new_account_id) = next_account_id { + *account_id = Some(new_account_id); + } + } + AuthData::ApiKey { .. } | AuthData::CodexAccessToken { .. } => { + anyhow::bail!("Cannot update OAuth tokens for this account type"); + } + } - let updated = update_account_chatgpt_tokens( - &account.id, - next_id_token, - refreshed.access_token, - next_refresh_token, - next_account_id, - claims.email, - claims.plan_type, - claims.subscription_expires_at, - )?; + if let Some(email) = claims.email { + stored.email = Some(email); + } + if let Some(plan_type) = claims.plan_type { + stored.plan_type = Some(plan_type); + } + if let Some(expires_at) = claims.subscription_expires_at { + stored.subscription_expires_at = Some(expires_at); + } - // Keep ~/.codex/auth.json in sync when this is the active account. - if is_active { - if let Err(err) = switch_to_account(&updated) { - println!("[Auth] Failed to sync active auth.json after token refresh: {err}"); + let updated = stored.clone(); + + // Keep ~/.codex/auth.json in sync when this is the active account. + // Checked and written under the store lock so a concurrent switch + // can't interleave and end up with auth.json holding stale tokens or + // the wrong account. Best-effort: the refreshed tokens must persist + // even if this sync write fails (the old refresh token may already + // be invalidated server-side). + if store.active_account_id.as_deref() == Some(updated.id.as_str()) { + if let Err(err) = switch_to_account(&updated) { + println!("[Auth] Failed to sync active auth.json after token refresh: {err}"); + } } - } - Ok(updated) + Ok(updated) + }) } /// Build a new ChatGPT account from a refresh token. diff --git a/src-tauri/src/commands/account.rs b/src-tauri/src/commands/account.rs index 5bc2e6e1..13781fc3 100644 --- a/src-tauri/src/commands/account.rs +++ b/src-tauri/src/commands/account.rs @@ -2,8 +2,8 @@ use crate::auth::{ add_account, create_chatgpt_account_from_refresh_token, get_active_account, - import_from_auth_json, import_from_auth_json_contents, load_accounts, remove_account, - save_accounts, set_active_account, switch_to_account, touch_account, + import_from_auth_json, import_from_auth_json_contents, load_accounts, mutate_accounts, + remove_account, switch_to_account, }; use crate::types::{AccountInfo, AccountsStore, AuthData, ImportAccountsSummary, StoredAccount}; @@ -158,25 +158,28 @@ pub async fn switch_account(account_id: String) -> Result<(), String> { } pub fn switch_account_by_id(account_id: &str) -> Result<(), String> { - let store = load_accounts().map_err(|e| e.to_string())?; - - // Find the account - let account = store - .accounts - .iter() - .find(|a| a.id == account_id) - .ok_or_else(|| format!("Account not found: {account_id}"))?; - ensure_codex_not_running()?; - // Write to ~/.codex/auth.json - switch_to_account(account).map_err(|e| e.to_string())?; + // Write auth.json and update the store in one locked cycle, so a + // concurrent switch (tray menu, LAN dashboard) can never leave auth.json + // pointing at one account while the store marks another as active. + mutate_accounts(|store| { + let account = store + .accounts + .iter() + .find(|a| a.id == account_id) + .cloned() + .with_context(|| format!("Account not found: {account_id}"))?; - // Update the active account in our store - set_active_account(account_id).map_err(|e| e.to_string())?; + switch_to_account(&account)?; - // Update last_used_at - touch_account(account_id).map_err(|e| e.to_string())?; + store.active_account_id = Some(account_id.to_string()); + if let Some(stored) = store.accounts.iter_mut().find(|a| a.id == account_id) { + stored.last_used_at = Some(chrono::Utc::now()); + } + Ok(()) + }) + .map_err(|e| e.to_string())?; // Restart Antigravity background process if it is running // This allows it to pick up the new authorization file seamlessly @@ -242,8 +245,12 @@ pub async fn import_accounts_slim_text(payload: String) -> Result Result { // Add the account to storage let stored = add_account(account).map_err(|e| e.to_string())?; - // Make it active and switch to it - set_active_account(&stored.id).map_err(|e| e.to_string())?; - switch_to_account(&stored).map_err(|e| e.to_string())?; - touch_account(&stored.id).map_err(|e| e.to_string())?; + // Make it active and switch to it in one locked cycle, so auth.json and + // the store's active account can't diverge under concurrent writers. + mutate_accounts(|store| { + switch_to_account(&stored)?; + store.active_account_id = Some(stored.id.clone()); + if let Some(account) = store.accounts.iter_mut().find(|a| a.id == stored.id) { + account.last_used_at = Some(chrono::Utc::now()); + } + Ok(()) + }) + .map_err(|e| e.to_string())?; let store = load_accounts().map_err(|e| e.to_string())?; let active_id = store.active_account_id.as_deref(); From c6d91963243cb75a4b4084ceca0bbd8d5abd049b Mon Sep 17 00:00:00 2001 From: zakyirsyaad Date: Thu, 9 Jul 2026 02:33:13 +0700 Subject: [PATCH 6/8] fix: pass PATH to Codex CLI child process --- AGENTS.md | 45 ++++++++++++++++++ src-tauri/src/auth/switcher.rs | 83 +++++++++++++++++++++++++++++++++- 2 files changed, 126 insertions(+), 2 deletions(-) create mode 100644 AGENTS.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..6a48d371 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,45 @@ +# Repository Guidelines + +## Project Structure & Module Organization + +Codex Switcher is a Tauri v2 desktop app with a React/TypeScript frontend and Rust backend. + +- `src/` contains the Vite/React UI. Key areas include `components/`, `hooks/`, `lib/`, and `types/`. +- `src-tauri/src/` contains Rust backend code. Commands live in `commands/`, account persistence and switching in `auth/`, API calls in `api/`, and tray/app menu logic in `tray.rs` and `app_menu.rs`. +- `src-tauri/icons/` stores app icons. Built frontend output goes to `dist/`. +- `scripts/` contains versioning, release, and Tauri wrapper scripts. +- Rust tests are colocated in `#[cfg(test)]` modules inside `src-tauri/src/**`. + +## Build, Test, and Development Commands + +- `pnpm install` installs frontend dependencies. +- `pnpm tauri dev` runs the desktop app in development mode. +- `pnpm build` runs TypeScript checking and builds the Vite frontend. +- `cargo test --manifest-path src-tauri/Cargo.toml` runs Rust unit tests. +- `pnpm lan` builds the frontend and serves the LAN dashboard through the `codex-web` binary. +- `pnpm tauri build` builds production bundles/installers. + +On Windows, use `pnpm tauri:win` instead of the POSIX wrapper. + +## Coding Style & Naming Conventions + +Use idiomatic Rust and TypeScript. Keep Rust modules focused by domain, and keep React components small enough to scan. Use `cargo fmt --manifest-path src-tauri/Cargo.toml` for Rust formatting. Frontend formatting follows the existing TypeScript/React style: 2-space indentation, PascalCase components, camelCase variables/functions, and explicit shared types in `src/types/`. + +## Testing Guidelines + +Prefer focused Rust unit tests near the code being changed. For storage, auth, token handling, and process logic, add regression tests for the exact failure mode. There is no frontend test runner configured; use `pnpm build` as the required frontend verification. For risky UI or app-flow changes, manually run `pnpm tauri dev`. + +## Commit & Pull Request Guidelines + +Git history uses short conventional-style messages, for example `fix: write auth files atomically` or `feat: add Codex access token accounts`. Keep commits scoped and descriptive. + +Pull requests should include: + +- Summary of behavior changed. +- Test plan with commands run. +- Screenshots for visible UI changes. +- Notes for auth, storage, updater, or release-impacting changes. + +## Security & Configuration Tips + +Never commit real tokens, account exports, `CODEX_ACCESS_TOKEN` values, or local credential files. `~/.codex-switcher/accounts.json` and `~/.codex/auth.json` contain secrets. When changing account persistence, preserve atomic writes and restrictive file permissions. diff --git a/src-tauri/src/auth/switcher.rs b/src-tauri/src/auth/switcher.rs index f92583ce..70d76189 100644 --- a/src-tauri/src/auth/switcher.rs +++ b/src-tauri/src/auth/switcher.rs @@ -115,7 +115,10 @@ fn login_with_codex_access_token(token: &str) -> Result<()> { anyhow::bail!("Codex access token is empty"); } - let mut child = Command::new(resolve_codex_executable()) + let codex_path = resolve_codex_executable(); + let child_path = child_path_for(&codex_path, std::env::var_os("PATH")); + let mut child = Command::new(&codex_path) + .env("PATH", child_path) .args(["login", "--with-access-token"]) .stdin(Stdio::piped()) .stdout(Stdio::piped()) @@ -176,6 +179,40 @@ fn resolve_codex_executable() -> PathBuf { PathBuf::from("codex") } +/// Build the PATH for the spawned Codex CLI child process. +/// +/// GUI apps launched from Finder/Dock inherit launchd's minimal PATH, and an +/// npm-installed `codex` is a `#!/usr/bin/env node` shim — if `node` isn't on +/// the *child's* PATH the login dies with "env: node: No such file or +/// directory" even though the shim itself was found. Prepend the resolved +/// codex binary's own directory (nvm/volta/npm colocate `node` with the shim) +/// plus the same well-known bin directories used to find codex itself. +fn child_path_for(codex_path: &Path, existing: Option) -> std::ffi::OsString { + let mut path_dirs: Vec = Vec::new(); + + if let Some(parent) = codex_path.parent() { + if !parent.as_os_str().is_empty() { + path_dirs.push(parent.to_path_buf()); + } + } + + #[cfg(unix)] + { + path_dirs.push(PathBuf::from("/opt/homebrew/bin")); + path_dirs.push(PathBuf::from("/usr/local/bin")); + if let Some(home) = dirs::home_dir() { + path_dirs.push(home.join(".local/bin")); + path_dirs.push(home.join(".npm-global/bin")); + } + } + + if let Some(existing) = existing.as_ref() { + path_dirs.extend(std::env::split_paths(existing)); + } + + std::env::join_paths(path_dirs).unwrap_or_else(|_| existing.unwrap_or_default()) +} + fn find_on_path(name: &str, path_var: Option) -> Option { let path_var = path_var?; std::env::split_paths(&path_var).find_map(|dir| { @@ -209,7 +246,49 @@ fn redact_access_token_from_output(output: &str, token: &str) -> String { #[cfg(test)] mod tests { - use super::{create_access_token_auth_json, find_on_path, redact_access_token_from_output}; + use super::{ + child_path_for, create_access_token_auth_json, find_on_path, + redact_access_token_from_output, + }; + + #[test] + fn child_path_puts_codex_parent_dir_first_and_keeps_existing_entries() { + let path = child_path_for( + std::path::Path::new("/some/toolchain/bin/codex"), + Some(std::ffi::OsString::from("/usr/bin:/bin")), + ); + let dirs: Vec = std::env::split_paths(&path).collect(); + + assert_eq!( + dirs.first(), + Some(&std::path::PathBuf::from("/some/toolchain/bin")), + "node must be findable next to the resolved codex shim" + ); + assert!(dirs.contains(&std::path::PathBuf::from("/usr/bin"))); + assert!(dirs.contains(&std::path::PathBuf::from("/bin"))); + } + + #[test] + fn child_path_adds_no_empty_entry_for_bare_codex_name() { + let path = child_path_for( + std::path::Path::new("codex"), + Some(std::ffi::OsString::from("/usr/bin")), + ); + let dirs: Vec = std::env::split_paths(&path).collect(); + + assert!(!dirs.iter().any(|d| d.as_os_str().is_empty())); + assert!(dirs.contains(&std::path::PathBuf::from("/usr/bin"))); + } + + #[cfg(unix)] + #[test] + fn child_path_includes_well_known_bin_dirs() { + let path = child_path_for(std::path::Path::new("codex"), None); + let dirs: Vec = std::env::split_paths(&path).collect(); + + assert!(dirs.contains(&std::path::PathBuf::from("/opt/homebrew/bin"))); + assert!(dirs.contains(&std::path::PathBuf::from("/usr/local/bin"))); + } #[test] fn find_on_path_returns_none_when_path_is_unset() { From de6c6a612887aa442d4ac667381adea1fca6e25e Mon Sep 17 00:00:00 2001 From: zakyirsyaad Date: Fri, 10 Jul 2026 17:57:06 +0700 Subject: [PATCH 7/8] fix: persist Codex access token auth directly --- src-tauri/src/auth/switcher.rs | 311 ++++++++------------------------- 1 file changed, 69 insertions(+), 242 deletions(-) diff --git a/src-tauri/src/auth/switcher.rs b/src-tauri/src/auth/switcher.rs index 70d76189..83802b1d 100644 --- a/src-tauri/src/auth/switcher.rs +++ b/src-tauri/src/auth/switcher.rs @@ -1,9 +1,7 @@ //! Account switching logic - writes credentials to ~/.codex/auth.json use std::fs; -use std::io::Write; use std::path::{Path, PathBuf}; -use std::process::{Command, Stdio}; use anyhow::{Context, Result}; use chrono::Utc; @@ -29,16 +27,23 @@ pub fn get_codex_auth_file() -> Result { Ok(get_codex_home()?.join("auth.json")) } -/// Switch to a specific account by writing its credentials to ~/.codex/auth.json +/// Switch to a specific account by writing its credentials to ~/.codex/auth.json. +/// +/// All auth modes — including `CodexAccessToken` — are materialized by writing +/// the same `auth.json` the official Codex CLI stores after a successful login. +/// We write it directly (atomically) rather than shelling out to +/// `codex login --with-access-token`: that command gates login on a network +/// fetch of the agent-identity JWKS which can hang or fail, and the file it +/// ultimately writes is byte-for-byte what `create_auth_json` produces here. pub fn switch_to_account(account: &StoredAccount) -> Result<()> { - if let AuthData::CodexAccessToken { token, .. } = &account.auth_data { - return login_with_codex_access_token(token); - } - let codex_home = get_codex_home()?; + write_account_auth_json(&codex_home, account) +} - // Ensure the codex home directory exists - fs::create_dir_all(&codex_home) +/// Write an account's credentials to `/auth.json` atomically, +/// creating the codex home directory if it does not exist. +fn write_account_auth_json(codex_home: &Path, account: &StoredAccount) -> Result<()> { + fs::create_dir_all(codex_home) .with_context(|| format!("Failed to create codex home: {}", codex_home.display()))?; let auth_json = create_auth_json(account)?; @@ -82,277 +87,92 @@ fn create_auth_json(account: &StoredAccount) -> Result { agent_identity: None, personal_access_token: None, }), - AuthData::CodexAccessToken { token, .. } => Ok(create_access_token_auth_json(token)), + AuthData::CodexAccessToken { token, .. } => create_access_token_auth_json(token), } } -fn create_access_token_auth_json(token: &str) -> AuthDotJson { - let trimmed = token.trim().to_string(); +fn create_access_token_auth_json(token: &str) -> Result { + let trimmed = token.trim(); + if trimmed.is_empty() { + anyhow::bail!("Codex access token is empty"); + } + if trimmed.starts_with("at-") { - AuthDotJson { + Ok(AuthDotJson { auth_mode: None, openai_api_key: None, tokens: None, last_refresh: None, agent_identity: None, - personal_access_token: Some(trimmed), - } + personal_access_token: Some(trimmed.to_string()), + }) } else { - AuthDotJson { + Ok(AuthDotJson { auth_mode: Some("agentIdentity".to_string()), openai_api_key: None, tokens: None, last_refresh: None, - agent_identity: Some(serde_json::Value::String(trimmed)), + agent_identity: Some(serde_json::Value::String(trimmed.to_string())), personal_access_token: None, - } - } -} - -fn login_with_codex_access_token(token: &str) -> Result<()> { - let trimmed = token.trim(); - if trimmed.is_empty() { - anyhow::bail!("Codex access token is empty"); - } - - let codex_path = resolve_codex_executable(); - let child_path = child_path_for(&codex_path, std::env::var_os("PATH")); - let mut child = Command::new(&codex_path) - .env("PATH", child_path) - .args(["login", "--with-access-token"]) - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .context("Failed to start Codex CLI. Make sure `codex` is installed and on PATH")?; - - { - let mut stdin = child - .stdin - .take() - .context("Failed to open Codex CLI stdin")?; - writeln!(stdin, "{trimmed}").context("Failed to send access token to Codex CLI")?; - } - - let output = child - .wait_with_output() - .context("Failed to wait for Codex CLI login")?; - - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); - let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string(); - let detail = if !stderr.is_empty() { - stderr - } else if !stdout.is_empty() { - stdout - } else { - format!("exit status {}", output.status) - }; - let detail = redact_access_token_from_output(&detail, trimmed); - anyhow::bail!("Codex access token login failed: {detail}"); - } - - Ok(()) -} - -/// Resolve the `codex` CLI executable, falling back to well-known install -/// locations when it isn't found on the current process's PATH. -/// -/// GUI apps launched from Finder/Dock/Spotlight on macOS inherit launchd's -/// minimal PATH, not the user's shell PATH — so a Homebrew-installed `codex` -/// (e.g. under `/opt/homebrew/bin`) is invisible to `Command::new("codex")` -/// even though it works fine from a terminal. -fn resolve_codex_executable() -> PathBuf { - if let Some(path) = find_on_path("codex", std::env::var_os("PATH")) { - return path; - } - - #[cfg(unix)] - { - for candidate in fallback_codex_paths(dirs::home_dir().as_deref()) { - if candidate.is_file() { - return candidate; - } - } - } - - PathBuf::from("codex") -} - -/// Build the PATH for the spawned Codex CLI child process. -/// -/// GUI apps launched from Finder/Dock inherit launchd's minimal PATH, and an -/// npm-installed `codex` is a `#!/usr/bin/env node` shim — if `node` isn't on -/// the *child's* PATH the login dies with "env: node: No such file or -/// directory" even though the shim itself was found. Prepend the resolved -/// codex binary's own directory (nvm/volta/npm colocate `node` with the shim) -/// plus the same well-known bin directories used to find codex itself. -fn child_path_for(codex_path: &Path, existing: Option) -> std::ffi::OsString { - let mut path_dirs: Vec = Vec::new(); - - if let Some(parent) = codex_path.parent() { - if !parent.as_os_str().is_empty() { - path_dirs.push(parent.to_path_buf()); - } - } - - #[cfg(unix)] - { - path_dirs.push(PathBuf::from("/opt/homebrew/bin")); - path_dirs.push(PathBuf::from("/usr/local/bin")); - if let Some(home) = dirs::home_dir() { - path_dirs.push(home.join(".local/bin")); - path_dirs.push(home.join(".npm-global/bin")); - } - } - - if let Some(existing) = existing.as_ref() { - path_dirs.extend(std::env::split_paths(existing)); - } - - std::env::join_paths(path_dirs).unwrap_or_else(|_| existing.unwrap_or_default()) -} - -fn find_on_path(name: &str, path_var: Option) -> Option { - let path_var = path_var?; - std::env::split_paths(&path_var).find_map(|dir| { - let candidate = dir.join(name); - candidate.is_file().then_some(candidate) - }) -} - -#[cfg(unix)] -fn fallback_codex_paths(home: Option<&Path>) -> Vec { - let mut candidates = vec![ - PathBuf::from("/opt/homebrew/bin/codex"), - PathBuf::from("/usr/local/bin/codex"), - ]; - - if let Some(home) = home { - candidates.push(home.join(".local/bin/codex")); - candidates.push(home.join(".npm-global/bin/codex")); + }) } - - candidates -} - -fn redact_access_token_from_output(output: &str, token: &str) -> String { - if token.is_empty() { - return output.to_string(); - } - - output.replace(token, "[redacted access token]") } #[cfg(test)] mod tests { - use super::{ - child_path_for, create_access_token_auth_json, find_on_path, - redact_access_token_from_output, - }; - - #[test] - fn child_path_puts_codex_parent_dir_first_and_keeps_existing_entries() { - let path = child_path_for( - std::path::Path::new("/some/toolchain/bin/codex"), - Some(std::ffi::OsString::from("/usr/bin:/bin")), - ); - let dirs: Vec = std::env::split_paths(&path).collect(); - - assert_eq!( - dirs.first(), - Some(&std::path::PathBuf::from("/some/toolchain/bin")), - "node must be findable next to the resolved codex shim" - ); - assert!(dirs.contains(&std::path::PathBuf::from("/usr/bin"))); - assert!(dirs.contains(&std::path::PathBuf::from("/bin"))); - } - - #[test] - fn child_path_adds_no_empty_entry_for_bare_codex_name() { - let path = child_path_for( - std::path::Path::new("codex"), - Some(std::ffi::OsString::from("/usr/bin")), - ); - let dirs: Vec = std::env::split_paths(&path).collect(); - - assert!(!dirs.iter().any(|d| d.as_os_str().is_empty())); - assert!(dirs.contains(&std::path::PathBuf::from("/usr/bin"))); - } - - #[cfg(unix)] - #[test] - fn child_path_includes_well_known_bin_dirs() { - let path = child_path_for(std::path::Path::new("codex"), None); - let dirs: Vec = std::env::split_paths(&path).collect(); + use super::{create_access_token_auth_json, write_account_auth_json}; + use crate::types::StoredAccount; - assert!(dirs.contains(&std::path::PathBuf::from("/opt/homebrew/bin"))); - assert!(dirs.contains(&std::path::PathBuf::from("/usr/local/bin"))); - } - - #[test] - fn find_on_path_returns_none_when_path_is_unset() { - assert_eq!(find_on_path("codex", None), None); - } - - #[test] - fn find_on_path_returns_none_when_not_present_in_any_directory() { - let path_var = std::ffi::OsString::from("/definitely/does/not/exist/xyz"); - assert_eq!(find_on_path("codex", Some(path_var)), None); - } - - #[test] - fn find_on_path_locates_executable_within_a_path_directory() { + fn unique_codex_home(label: &str) -> std::path::PathBuf { let dir = std::env::temp_dir().join(format!( - "codex-switcher-find-on-path-test-{}", - std::process::id() + "codex-switcher-switch-test-{}-{}", + std::process::id(), + label )); - std::fs::create_dir_all(&dir).unwrap(); - let exe = dir.join("codex"); - std::fs::write(&exe, b"").unwrap(); - - let path_var = std::ffi::OsString::from(dir.as_os_str()); - let found = find_on_path("codex", Some(path_var)); - std::fs::remove_dir_all(&dir).ok(); - - assert_eq!(found.as_deref(), Some(exe.as_path())); + dir } - #[cfg(unix)] #[test] - fn fallback_codex_paths_includes_common_homebrew_prefixes() { - let candidates = super::fallback_codex_paths(None); - - assert!(candidates.contains(&std::path::PathBuf::from("/opt/homebrew/bin/codex"))); - assert!(candidates.contains(&std::path::PathBuf::from("/usr/local/bin/codex"))); + fn switching_jwt_access_token_account_writes_agent_identity_auth_json() { + // A successful `codex login --with-access-token` writes exactly this + // file; we produce it directly instead of shelling out to the CLI (whose + // login gates on an agent-identity JWKS fetch that can hang/fail). + let token = ["header", "payload", "signature"].join("."); + let account = StoredAccount::new_codex_access_token("K12".to_string(), token.clone()); + let codex_home = unique_codex_home("jwt"); + + write_account_auth_json(&codex_home, &account).expect("switch should write auth.json"); + + let written = std::fs::read_to_string(codex_home.join("auth.json")).unwrap(); + let parsed: serde_json::Value = serde_json::from_str(&written).unwrap(); + std::fs::remove_dir_all(&codex_home).ok(); + + assert_eq!(parsed["auth_mode"], "agentIdentity"); + assert_eq!(parsed["agent_identity"], token); + assert!(parsed.get("tokens").is_none()); } - #[cfg(unix)] #[test] - fn fallback_codex_paths_includes_home_relative_locations_when_home_is_known() { - let home = std::path::PathBuf::from("/Users/example"); - let candidates = super::fallback_codex_paths(Some(&home)); + fn switching_personal_access_token_account_writes_personal_access_token() { + let account = + StoredAccount::new_codex_access_token("K12".to_string(), "at-secret-123".to_string()); + let codex_home = unique_codex_home("pat"); - assert!(candidates.contains(&home.join(".local/bin/codex"))); - assert!(candidates.contains(&home.join(".npm-global/bin/codex"))); - } + write_account_auth_json(&codex_home, &account).expect("switch should write auth.json"); - #[test] - fn redacts_access_token_from_cli_error_output() { - let marker = ["sample", "token", "value"].join("-"); - let output = format!("login failed for {marker}"); + let written = std::fs::read_to_string(codex_home.join("auth.json")).unwrap(); + let parsed: serde_json::Value = serde_json::from_str(&written).unwrap(); + std::fs::remove_dir_all(&codex_home).ok(); - assert_eq!( - redact_access_token_from_output(&output, &marker), - "login failed for [redacted access token]" - ); + assert_eq!(parsed["personal_access_token"], "at-secret-123"); + assert!(parsed.get("auth_mode").is_none()); } #[test] fn creates_agent_identity_auth_json_for_codex_access_token_jwt() { let sample_access_token = ["header", "payload", "signature"].join("."); - let auth = create_access_token_auth_json(&sample_access_token); + let auth = create_access_token_auth_json(&sample_access_token).unwrap(); assert_eq!(auth.auth_mode.as_deref(), Some("agentIdentity")); assert_eq!( @@ -364,6 +184,13 @@ mod tests { assert!(auth.tokens.is_none()); assert!(auth.personal_access_token.is_none()); } + + #[test] + fn rejects_empty_codex_access_token() { + let error = create_access_token_auth_json(" \n\t ").unwrap_err(); + + assert!(error.to_string().contains("access token is empty")); + } } /// Import an account from an existing auth.json file From ba00213c072e87d941793c890ad355b9c257196a Mon Sep 17 00:00:00 2001 From: Nant361 Date: Sun, 12 Jul 2026 20:41:00 +0700 Subject: [PATCH 8/8] feat: support K12 access token usage stats --- src-tauri/src/api/usage.rs | 2 +- src-tauri/src/commands/account_stats.rs | 133 +++++++++++++++++++----- src/components/AccountCard.tsx | 4 +- src/components/AccountUsageStats.tsx | 2 +- 4 files changed, 112 insertions(+), 29 deletions(-) diff --git a/src-tauri/src/api/usage.rs b/src-tauri/src/api/usage.rs index 90800251..6a1e6835 100644 --- a/src-tauri/src/api/usage.rs +++ b/src-tauri/src/api/usage.rs @@ -362,7 +362,7 @@ fn build_chatgpt_headers( Ok(headers) } -async fn build_codex_access_token_headers(account: &StoredAccount) -> Result { +pub(crate) async fn build_codex_access_token_headers(account: &StoredAccount) -> Result { let AuthData::CodexAccessToken { token, account_id, diff --git a/src-tauri/src/commands/account_stats.rs b/src-tauri/src/commands/account_stats.rs index d18b7ac5..5461a798 100644 --- a/src-tauri/src/commands/account_stats.rs +++ b/src-tauri/src/commands/account_stats.rs @@ -7,6 +7,7 @@ use reqwest::{ }; use serde::{Deserialize, Serialize}; +use crate::api::build_codex_access_token_headers; use crate::auth::{ensure_chatgpt_tokens_fresh, load_accounts, refresh_chatgpt_tokens}; use crate::types::{AuthData, AuthMode, StoredAccount}; @@ -192,11 +193,8 @@ pub async fn get_account_usage_stats(account_id: String) -> Result Result bool { + matches!(auth_mode, AuthMode::ChatGPT | AuthMode::CodexAccessToken) +} + +fn api_key_unavailable_message() -> &'static str { + "Usage stats are unavailable for API key accounts." +} + async fn fetch_profile_usage(account: &StoredAccount) -> anyhow::Result { + match account.auth_mode { + AuthMode::ChatGPT => fetch_chatgpt_profile_usage(account).await, + AuthMode::CodexAccessToken => fetch_codex_access_token_profile_usage(account).await, + AuthMode::ApiKey => Ok(unavailable_stats( + account.id.clone(), + api_key_unavailable_message(), + )), + } +} + +async fn fetch_chatgpt_profile_usage(account: &StoredAccount) -> anyhow::Result { let fresh_account = ensure_chatgpt_tokens_fresh(account).await?; - let mut response = send_profile_usage_request(&fresh_account).await?; + let mut headers = chatgpt_headers_for_account(&fresh_account)?; + let mut response = send_profile_usage_request(headers.clone()).await?; if response.status() == StatusCode::UNAUTHORIZED { let refreshed_account = refresh_chatgpt_tokens(&fresh_account).await?; - response = send_profile_usage_request(&refreshed_account).await?; - return parse_profile_usage_with_reset_credits(&refreshed_account, response).await; + headers = chatgpt_headers_for_account(&refreshed_account)?; + response = send_profile_usage_request(headers.clone()).await?; + return parse_profile_usage_with_reset_credits(&refreshed_account.id, response, headers) + .await; } - parse_profile_usage_with_reset_credits(&fresh_account, response).await + parse_profile_usage_with_reset_credits(&fresh_account.id, response, headers).await } -async fn send_profile_usage_request(account: &StoredAccount) -> anyhow::Result { +async fn fetch_codex_access_token_profile_usage( + account: &StoredAccount, +) -> anyhow::Result { + let headers = build_codex_access_token_headers(account).await?; + let response = send_profile_usage_request(headers.clone()).await?; + + parse_profile_usage_with_reset_credits(&account.id, response, headers).await +} + +fn chatgpt_headers_for_account(account: &StoredAccount) -> anyhow::Result { let (access_token, chatgpt_account_id) = extract_chatgpt_auth(account)?; + build_chatgpt_headers(access_token, chatgpt_account_id) +} + +async fn send_profile_usage_request(headers: HeaderMap) -> anyhow::Result { let client = reqwest::Client::new(); Ok(client .get(CHATGPT_PROFILE_USAGE_URL) - .headers(build_chatgpt_headers(access_token, chatgpt_account_id)?) + .headers(headers) .send() .await?) } async fn parse_profile_usage_with_reset_credits( - account: &StoredAccount, + account_id: &str, response: reqwest::Response, + headers: HeaderMap, ) -> anyhow::Result { - let mut stats = parse_profile_usage_response(&account.id, response).await?; + let mut stats = parse_profile_usage_response(account_id, response).await?; if stats.available { - stats.reset_credits = fetch_reset_credits(account).await.ok(); + stats.reset_credits = fetch_reset_credits(headers).await.ok(); } Ok(stats) @@ -257,15 +291,11 @@ async fn parse_profile_usage_response( Ok(map_profile_usage(account_id, payload)) } -async fn fetch_reset_credits(account: &StoredAccount) -> anyhow::Result { - let (access_token, chatgpt_account_id) = extract_chatgpt_auth(account)?; +async fn fetch_reset_credits(headers: HeaderMap) -> anyhow::Result { let client = reqwest::Client::new(); let response = client .get(CHATGPT_RESET_CREDITS_URL) - .headers(build_reset_credits_headers( - access_token, - chatgpt_account_id, - )?) + .headers(build_reset_credits_headers(headers)) .send() .await?; @@ -421,11 +451,7 @@ fn build_chatgpt_headers( Ok(headers) } -fn build_reset_credits_headers( - access_token: &str, - chatgpt_account_id: Option<&str>, -) -> anyhow::Result { - let mut headers = build_chatgpt_headers(access_token, chatgpt_account_id)?; +fn build_reset_credits_headers(mut headers: HeaderMap) -> HeaderMap { headers.insert(ACCEPT, HeaderValue::from_static("application/json")); headers.insert( HeaderName::from_static("openai-beta"), @@ -435,7 +461,7 @@ fn build_reset_credits_headers( HeaderName::from_static("originator"), HeaderValue::from_static("Codex Desktop"), ); - Ok(headers) + headers } fn extract_chatgpt_auth(account: &StoredAccount) -> anyhow::Result<(&str, Option<&str>)> { @@ -455,6 +481,61 @@ fn extract_chatgpt_auth(account: &StoredAccount) -> anyhow::Result<(&str, Option mod tests { use super::*; + #[test] + fn usage_stats_support_chatgpt_and_codex_access_tokens() { + assert!(supports_usage_stats(AuthMode::ChatGPT)); + assert!(supports_usage_stats(AuthMode::CodexAccessToken)); + assert!(!supports_usage_stats(AuthMode::ApiKey)); + assert_eq!( + api_key_unavailable_message(), + "Usage stats are unavailable for API key accounts." + ); + } + + #[test] + fn reset_credit_headers_preserve_agent_identity_authentication() { + let mut headers = HeaderMap::new(); + headers.insert( + AUTHORIZATION, + HeaderValue::from_static("AgentAssertion signed-assertion"), + ); + headers.insert( + HeaderName::from_static("chatgpt-account-id"), + HeaderValue::from_static("account-123"), + ); + + let headers = build_reset_credits_headers(headers); + + assert_eq!( + headers + .get(AUTHORIZATION) + .and_then(|value| value.to_str().ok()), + Some("AgentAssertion signed-assertion") + ); + assert_eq!( + headers + .get("chatgpt-account-id") + .and_then(|value| value.to_str().ok()), + Some("account-123") + ); + assert_eq!( + headers.get(ACCEPT).and_then(|value| value.to_str().ok()), + Some("application/json") + ); + assert_eq!( + headers + .get("openai-beta") + .and_then(|value| value.to_str().ok()), + Some("codex-1") + ); + assert_eq!( + headers + .get("originator") + .and_then(|value| value.to_str().ok()), + Some("Codex Desktop") + ); + } + #[test] fn profile_usage_response_maps_profile_stats() { let payload: ProfileUsageResponse = serde_json::from_value(serde_json::json!({ diff --git a/src/components/AccountCard.tsx b/src/components/AccountCard.tsx index dd30ae45..124ca997 100644 --- a/src/components/AccountCard.tsx +++ b/src/components/AccountCard.tsx @@ -242,6 +242,8 @@ export function AccountCard({ const planKey = account.plan_type?.toLowerCase() || account.auth_mode; const planColorClass = planColors[planKey] || planColors.free; + const supportsUsageStats = + account.auth_mode === "chat_g_p_t" || account.auth_mode === "codex_access_token"; const showSubscriptionStatus = account.auth_mode === "chat_g_p_t"; const subscriptionStatus = getSubscriptionStatus(account.subscription_expires_at); const resetCreditsCount = formatResetCreditsCount(resetCredits); @@ -417,7 +419,7 @@ export function AccountCard({ diff --git a/src/components/AccountUsageStats.tsx b/src/components/AccountUsageStats.tsx index 3daedf2e..6c909164 100644 --- a/src/components/AccountUsageStats.tsx +++ b/src/components/AccountUsageStats.tsx @@ -356,7 +356,7 @@ export function AccountUsageStats({ const requestId = ++requestSeq.current; if (!enabled) { - const next = emptyStats(accountId, "Usage stats are available for ChatGPT accounts only."); + const next = emptyStats(accountId, "Usage stats are unavailable for API key accounts."); setStats(next); onStatsLoaded?.(next); setLoading(false);