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/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 8c304b64..959b6cf8 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", @@ -576,6 +593,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" @@ -686,6 +709,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" @@ -736,6 +789,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" @@ -770,6 +850,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" @@ -954,6 +1044,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" @@ -1073,6 +1187,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" @@ -1406,6 +1526,7 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", + "zeroize", ] [[package]] @@ -3019,6 +3140,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" @@ -3644,6 +3775,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" @@ -3995,6 +4135,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" @@ -4099,6 +4248,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" @@ -5050,9 +5209,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 92cf34bb..668f1434 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -39,6 +39,8 @@ url = "2" flate2 = "1" chacha20poly1305 = "0.10" pbkdf2 = "0.12" +ed25519-dalek = { version = "2.2", features = ["pkcs8"] } +crypto_box = { version = "0.9", features = ["seal"] } [target.'cfg(target_os = "macos")'.dependencies] plist = "1" diff --git a/src-tauri/src/api/usage.rs b/src-tauri/src/api/usage.rs index 6c4c1331..6a1e6835 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) } +pub(crate) 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..6ec1f356 100644 --- a/src-tauri/src/auth/storage.rs +++ b/src-tauri/src/auth/storage.rs @@ -1,12 +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, StoredAccount}; -use crate::types::{AccountsStore, AppSettings, AuthData, 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 { @@ -16,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)) +} + +/// 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 +/// 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()); @@ -34,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()))?; @@ -50,114 +99,192 @@ 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")?; - fs::write(&path, content) - .with_context(|| format!("Failed to write settings file: {}", path.display()))?; - - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let perms = fs::Permissions::from_mode(0o600); - fs::set_permissions(&path, perms)?; - } + 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) } -/// Save the accounts store to disk -pub fn save_accounts(store: &AccountsStore) -> Result<()> { - let path = get_accounts_file()?; +fn mutate_accounts_in( + dir: &Path, + mutate: impl FnOnce(&mut AccountsStore) -> Result, +) -> Result { + let _lock = acquire_store_lock(dir)?; - // 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()))?; - } + 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")?; + 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 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. 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())); + + let result = write_temp_then_rename(&tmp_path, path, content); + if result.is_err() { + let _ = fs::remove_file(&tmp_path); + } + result +} + +/// 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."); - fs::write(&path, content) - .with_context(|| format!("Failed to write accounts file: {}", path.display()))?; + 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()); + } + } +} - // Set restrictive permissions on Unix +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; - let perms = fs::Permissions::from_mode(0o600); - fs::set_permissions(&path, perms)?; + use std::os::unix::fs::OpenOptionsExt; + options.mode(0o600); } + 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(()) } -/// Add a new account to the store -pub fn add_account(account: StoredAccount) -> Result { - let mut store = load_accounts()?; +#[cfg(not(windows))] +fn rename_replacing(from: &Path, to: &Path) -> std::io::Result<()> { + fs::rename(from, to) +} - // Check for duplicate names - if store.accounts.iter().any(|a| a.name == account.name) { - anyhow::bail!("An account with name '{}' already exists", account.name); +/// 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), + } } + Err(last_error.expect("rename attempted at least once")) +} - let account_clone = account.clone(); - store.accounts.push(account); +/// Add a new account to the store +pub fn add_account(account: StoredAccount) -> Result { + mutate_accounts(|store| { + if store.accounts.iter().any(|a| a.name == account.name) { + anyhow::bail!("An account with name '{}' already exists", account.name); + } - // If this is the first account, make it active - if store.accounts.len() == 1 { - store.active_account_id = Some(account_clone.id.clone()); - } + 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()); + } - 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}"); - } + mutate_accounts(|store| { + let initial_len = store.accounts.len(); + store.accounts.retain(|a| a.id != 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()); - } - - 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 @@ -176,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, @@ -196,112 +311,232 @@ pub fn update_account_metadata( plan_type: Option, subscription_expires_at: Option>>, ) -> Result { - let mut store = load_accounts()?; + 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"); + } + } - // Check for duplicate names first (if renaming) - if let Some(ref new_name) = name { - if store + // Now find and update the account + let account = store .accounts - .iter() - .any(|a| a.id != account_id && a.name == *new_name) - { - anyhow::bail!("An account with name '{new_name}' already exists"); + .iter_mut() + .find(|a| a.id == account_id) + .context("Account not found")?; + + if let Some(new_name) = name { + account.name = new_name; } - } - // Now find and update the account - let account = store - .accounts - .iter_mut() - .find(|a| a.id == account_id) - .context("Account not found")?; + if email.is_some() { + account.email = email; + } - if let Some(new_name) = name { - account.name = new_name; - } + if plan_type.is_some() { + account.plan_type = plan_type; + } - if email.is_some() { - account.email = email; - } + if let Some(subscription_expires_at) = subscription_expires_at { + account.subscription_expires_at = subscription_expires_at; + } - if plan_type.is_some() { - account.plan_type = plan_type; - } + Ok(account.clone()) + }) +} - if let Some(subscription_expires_at) = subscription_expires_at { - account.subscription_expires_at = subscription_expires_at; - } +/// Get the list of masked account IDs +pub fn get_masked_account_ids() -> Result> { + let store = load_accounts()?; + Ok(store.masked_account_ids.clone()) +} - let updated = account.clone(); - save_accounts(&store)?; - Ok(updated) +/// Set the list of masked account IDs +pub fn set_masked_account_ids(ids: Vec) -> Result<()> { + mutate_accounts(|store| { + store.masked_account_ids = ids; + Ok(()) + }) } -/// 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); +#[cfg(test)] +mod tests { + use super::*; + 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 = temp_store_dir(); + 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; } } - AuthData::ApiKey { .. } => { - anyhow::bail!("Cannot update OAuth tokens for an API key account"); + writer.join().unwrap(); + assert!( + reads > 0, + "test didn't overlap reads with writes — increase iterations" + ); + + 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); } - if let Some(new_email) = email { - account.email = Some(new_email); + #[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); } - if let Some(new_plan_type) = plan_type { - account.plan_type = Some(new_plan_type); + #[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); } - if let Some(subscription_expires_at) = subscription_expires_at { - account.subscription_expires_at = Some(subscription_expires_at); + #[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); } - let updated = account.clone(); - save_accounts(&store)?; - Ok(updated) -} + #[test] + fn sweeps_stale_tmp_files_and_leaves_none_behind() { + let dir = temp_store_dir(); + let path = dir.join("accounts.json"); -/// Get the list of masked account IDs -pub fn get_masked_account_ids() -> Result> { - let store = load_accounts()?; - Ok(store.masked_account_ids.clone()) -} + fs::write(dir.join("accounts.json.tmp.stale-crash-leftover"), "junk").unwrap(); + fs::write(dir.join("unrelated.txt"), "keep me").unwrap(); -/// 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(()) + 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/switcher.rs b/src-tauri/src/auth/switcher.rs index e213574e..83802b1d 100644 --- a/src-tauri/src/auth/switcher.rs +++ b/src-tauri/src/auth/switcher.rs @@ -1,11 +1,12 @@ //! Account switching logic - writes credentials to ~/.codex/auth.json use std::fs; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; 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, }; @@ -26,12 +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<()> { 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)?; @@ -40,17 +52,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(()) } @@ -58,9 +62,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 +75,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 +84,112 @@ 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, .. } => create_access_token_auth_json(token), + } +} + +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-") { + Ok(AuthDotJson { + auth_mode: None, + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: Some(trimmed.to_string()), + }) + } else { + 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.to_string())), + personal_access_token: None, + }) + } +} + +#[cfg(test)] +mod tests { + use super::{create_access_token_auth_json, write_account_auth_json}; + use crate::types::StoredAccount; + + fn unique_codex_home(label: &str) -> std::path::PathBuf { + let dir = std::env::temp_dir().join(format!( + "codex-switcher-switch-test-{}-{}", + std::process::id(), + label + )); + std::fs::remove_dir_all(&dir).ok(); + dir + } + + #[test] + 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()); + } + + #[test] + 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"); + + 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["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).unwrap(); + + 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()); + } + + #[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")); } } @@ -113,8 +226,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 +263,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..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"; @@ -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, @@ -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 915bb77b..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}; @@ -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> { @@ -132,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 @@ -216,8 +245,12 @@ pub async fn import_accounts_slim_text(payload: String) -> Result 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..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>)> { @@ -445,7 +471,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") + } } } @@ -453,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-tauri/src/commands/oauth.rs b/src-tauri/src/commands/oauth.rs index f4918147..cd005f18 100644 --- a/src-tauri/src/commands/oauth.rs +++ b/src-tauri/src/commands/oauth.rs @@ -5,9 +5,7 @@ use std::sync::{Arc, Mutex}; use tokio::sync::oneshot; use crate::auth::oauth_server::{start_oauth_login, wait_for_oauth_login, OAuthLoginResult}; -use crate::auth::{ - add_account, load_accounts, set_active_account, switch_to_account, touch_account, -}; +use crate::auth::{add_account, load_accounts, mutate_accounts, switch_to_account}; use crate::types::{AccountInfo, OAuthLoginInfo}; struct PendingOAuth { @@ -59,10 +57,17 @@ pub async fn complete_login() -> 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(); 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..46f562c9 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -164,6 +164,7 @@ function App() { deleteAccount, renameAccount, importFromFile, + addFromAccessToken, exportAccountsSlimText, importAccountsSlimText, startOAuthLogin, @@ -501,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); } @@ -1825,6 +1828,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..124ca997 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,10 +237,13 @@ 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 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); @@ -414,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); diff --git a/src/components/AddAccountModal.tsx b/src/components/AddAccountModal.tsx index 9b795dec..f4b1acc6 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 K12 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 K12 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,28 @@ 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 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" + /> +
+ )} + {/* Error */} {error && (
@@ -267,15 +340,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 {