diff --git a/AGENTS.md b/AGENTS.md index 3a287a599..8c63fcf36 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -20,10 +20,6 @@ In the codex-rs folder where the rust code lives: - When writing tests, prefer comparing the equality of entire objects over fields one by one. - When making a change that adds or changes an API, ensure that the documentation in the `docs/` folder is up to date if applicable. - If you change `ConfigToml` or nested config types, run `just write-config-schema` to update `codex-rs/core/config.schema.json`. -- If you change Rust dependencies (`Cargo.toml` or `Cargo.lock`), run `just bazel-lock-update` from the - repo root to refresh `MODULE.bazel.lock`, and include that lockfile update in the same change. -- After dependency changes, run `just bazel-lock-check` from the repo root so lockfile drift is caught - locally before CI. - Bazel does not automatically make source-tree files available to compile-time Rust file access. If you add `include_str!`, `include_bytes!`, `sqlx::migrate!`, or similar build-time file or directory reads, update the crate's `BUILD.bazel` (`compile_data`, `build_script_data`, or test diff --git a/README.md b/README.md index 1e44875f2..1487e285c 100644 --- a/README.md +++ b/README.md @@ -1,60 +1,128 @@ -

npm i -g @openai/codex
or brew install --cask codex

-

Codex CLI is a coding agent from OpenAI that runs locally on your computer. -

- Codex CLI splash -

-
-If you want Codex in your code editor (VS Code, Cursor, Windsurf), install in your IDE. -
If you want the desktop app experience, run codex app or visit the Codex App page. -
If you are looking for the cloud-based agent from OpenAI, Codex Web, go to chatgpt.com/codex.

- ---- - -## Quickstart - -### Installing and running Codex CLI - -Install globally with your preferred package manager: - -```shell -# Install using npm -npm install -g @openai/codex +# Codex Enhanced + +Codex Enhanced is a standalone public distribution of Codex focused on +multi-account ChatGPT operations, fullscreen TUI workflow improvements, and a +smaller long-term fork maintenance surface. + +This repository is maintained as its own GitHub project instead of a GitHub +fork. It tracks upstream Codex where practical, while keeping product-specific +behavior behind a dedicated extension layer so future changes can converge on +plugins instead of repeated invasive rebases. + +## Why This Version Exists + +Upstream Codex already provides a strong local coding agent. The main problem +this project solves is operational: + +- switching between multiple ChatGPT accounts should not require manual file + juggling +- rate-limit and usage-limit handling should be able to fail over to another + account automatically +- fullscreen TUI workflows should expose session and account operations in one + operator-facing control surface +- fork-specific behavior should move toward a stable extension boundary instead + of expanding the patch set in core runtime code + +## What Is Included + +- Managed account storage under `~/.codex/accounts` +- `codex login --auth ` for capturing multiple ChatGPT logins into named + account slots +- Account pool metadata with stable IDs, aliases, cooldown state, and inferred + usage windows +- Threshold-based account routing and one-shot retry on explicit + limit/rejection failures for normal user turns in the fullscreen TUI path +- A `Ctrl-P` control panel with: + - global session picker + - account selection + - alias rename submenu + - current-session fork entry point +- A dedicated `codex-rs/ext` crate for fork-owned extension state and host + compatibility groundwork + +## Current Scope + +The current milestone is a practical MVP for daily use, not the final extension +architecture. + +Implemented now: + +- managed ChatGPT account registry and auth snapshot layout +- account activation and alias management in the TUI +- control-panel-driven session and account operations +- inferred cooldown recording from explicit limit errors +- login-time account registration + +Planned next: + +- broader automatic account routing coverage beyond the current fullscreen TUI + path +- observable switch reasons and richer operator status views +- hook/interceptor expansion +- capability-negotiated WASM plugins built on top of `codex-ext` + +## Repository Layout + +- [codex-rs/ext](./codex-rs/ext) + Fork-owned extension crate for account pool state, auth snapshots, and future + plugin host compatibility. +- [codex-rs/tui](./codex-rs/tui) + Fullscreen local TUI implementation. +- [codex-rs/tui_app_server](./codex-rs/tui_app_server) + App-server-backed TUI implementation that mirrors relevant UX changes. +- [docs/fork-extension-mvp.md](./docs/fork-extension-mvp.md) + Fork proposal, MVP design, and phased roadmap. + +## Build And Run + +Build the Rust CLI locally: + +```bash +cd codex-rs +cargo build -p codex-cli +./target/debug/codex ``` -```shell -# Install using Homebrew -brew install --cask codex -``` +Install it into your shell path: -Then simply run `codex` to get started. +```bash +cd codex-rs +cargo build --release -p codex-cli +sudo ln -sf "$(pwd)/target/release/codex" /usr/local/bin/codex +codex --help +``` -
-You can also go to the latest GitHub Release and download the appropriate binary for your platform. +## Managed Account Quick Start -Each GitHub Release contains many executables, but in practice, you likely want one of these: +Register multiple ChatGPT logins into the managed account pool: -- macOS - - Apple Silicon/arm64: `codex-aarch64-apple-darwin.tar.gz` - - x86_64 (older Mac hardware): `codex-x86_64-apple-darwin.tar.gz` -- Linux - - x86_64: `codex-x86_64-unknown-linux-musl.tar.gz` - - arm64: `codex-aarch64-unknown-linux-musl.tar.gz` +```bash +codex login --auth primary +codex login --auth backup +codex login status +``` -Each archive contains a single entry with the platform baked into the name (e.g., `codex-x86_64-unknown-linux-musl`), so you likely want to rename it to `codex` after extracting it. +Start Codex, then use: -
+- `Ctrl-P -> Sessions` to open the global session picker +- `Ctrl-P -> Accounts` to switch the active managed account +- `Ctrl-P -> Accounts -> Rename` to rename account aliases -### Using Codex with your ChatGPT plan +Managed account state is stored under: -Run `codex` and select **Sign in with ChatGPT**. We recommend signing into your ChatGPT account to use Codex as part of your Plus, Pro, Team, Edu, or Enterprise plan. [Learn more about what's included in your ChatGPT plan](https://help.openai.com/en/articles/11369540-codex-in-chatgpt). +```text +~/.codex/accounts/ +├── account-pool.json +└── / + └── auth.json +``` -You can also use Codex with an API key, but this requires [additional setup](https://developers.openai.com/codex/auth#sign-in-with-an-api-key). +## Upstream Relationship -## Docs +This project is based on OpenAI Codex and keeps upstream history so changes can +be rebased and audited cleanly. The maintenance goal is to keep the fork-owned +delta small, explicit, and increasingly isolated behind `codex-ext`. -- [**Codex Documentation**](https://developers.openai.com/codex) -- [**Contributing**](./docs/contributing.md) -- [**Installing & building**](./docs/install.md) -- [**Open source fund**](./docs/open-source-fund.md) +## License -This repository is licensed under the [Apache-2.0 License](LICENSE). +This repository remains licensed under the [Apache-2.0 License](LICENSE). diff --git a/codex-cli/Dockerfile b/codex-cli/Dockerfile index 21a90a483..8cdb3dc90 100644 --- a/codex-cli/Dockerfile +++ b/codex-cli/Dockerfile @@ -1,4 +1,4 @@ -FROM node:24-slim +FROM node:25-slim ARG TZ ENV TZ="$TZ" diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d0917ee38..e1bd42d14 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1660,6 +1660,7 @@ dependencies = [ "anyhow", "assert_cmd", "assert_matches", + "base64 0.22.1", "clap", "clap_complete", "codex-app-server", @@ -1672,6 +1673,7 @@ dependencies = [ "codex-core", "codex-exec", "codex-execpolicy", + "codex-ext", "codex-features", "codex-login", "codex-mcp-server", @@ -2092,6 +2094,19 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "codex-ext" +version = "0.0.0" +dependencies = [ + "base64 0.22.1", + "codex-app-server-protocol", + "codex-core", + "pretty_assertions", + "serde", + "serde_json", + "tempfile", +] + [[package]] name = "codex-features" version = "0.0.0" @@ -2602,6 +2617,7 @@ dependencies = [ "codex-client", "codex-cloud-requirements", "codex-core", + "codex-ext", "codex-features", "codex-feedback", "codex-file-search", @@ -2695,6 +2711,7 @@ dependencies = [ "codex-client", "codex-cloud-requirements", "codex-core", + "codex-ext", "codex-features", "codex-feedback", "codex-file-search", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 524b61e3b..dcaf32e1f 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -27,6 +27,7 @@ members = [ "hooks", "secrets", "exec", + "ext", "exec-server", "execpolicy", "execpolicy-legacy", @@ -110,6 +111,7 @@ codex-connectors = { path = "connectors" } codex-config = { path = "config" } codex-core = { path = "core" } codex-exec = { path = "exec" } +codex-ext = { path = "ext" } codex-exec-server = { path = "exec-server" } codex-execpolicy = { path = "execpolicy" } codex-experimental-api-macros = { path = "codex-experimental-api-macros" } diff --git a/codex-rs/cli/Cargo.toml b/codex-rs/cli/Cargo.toml index c2fd1300c..028a8991d 100644 --- a/codex-rs/cli/Cargo.toml +++ b/codex-rs/cli/Cargo.toml @@ -30,6 +30,7 @@ codex-config = { workspace = true } codex-core = { workspace = true } codex-exec = { workspace = true } codex-execpolicy = { workspace = true } +codex-ext = { workspace = true } codex-features = { workspace = true } codex-login = { workspace = true } codex-mcp-server = { workspace = true } @@ -65,6 +66,7 @@ codex_windows_sandbox = { package = "codex-windows-sandbox", path = "../windows- [dev-dependencies] assert_cmd = { workspace = true } assert_matches = { workspace = true } +base64 = { workspace = true } codex-utils-cargo-bin = { workspace = true } predicates = { workspace = true } pretty_assertions = { workspace = true } diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index d0cc1a3a1..1e9ae03c4 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -11,9 +11,14 @@ use codex_core::CodexAuth; use codex_core::auth::AuthCredentialsStoreMode; use codex_core::auth::AuthMode; use codex_core::auth::CLIENT_ID; +use codex_core::auth::load_auth_dot_json; use codex_core::auth::login_with_api_key; use codex_core::auth::logout; use codex_core::config::Config; +use codex_ext::AccountPoolStore; +use codex_ext::ManagedAccountSnapshot; +use codex_ext::persist_current_managed_account_snapshot; +use codex_ext::persist_managed_account_auth_snapshot; use codex_login::ServerOptions; use codex_login::run_device_code_login; use codex_login::run_login_server; @@ -22,7 +27,10 @@ use codex_utils_cli::CliConfigOverrides; use std::fs::OpenOptions; use std::io::IsTerminal; use std::io::Read; +use std::path::Path; use std::path::PathBuf; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; use tracing_appender::non_blocking; use tracing_appender::non_blocking::WorkerGuard; use tracing_subscriber::EnvFilter; @@ -110,13 +118,67 @@ fn print_login_server_start(actual_port: u16, auth_url: &str) { ); } +fn now_timestamp() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .ok() + .and_then(|duration| i64::try_from(duration.as_secs()).ok()) + .unwrap_or(i64::MAX) +} + +fn upsert_managed_account_snapshot( + codex_home: &Path, + snapshot: &ManagedAccountSnapshot, + set_active: bool, +) -> std::io::Result<()> { + let account_id = snapshot.profile.id.clone(); + AccountPoolStore::new(codex_home.to_path_buf()).update(|state| { + state.upsert_account(snapshot.profile.clone()); + if set_active { + state.set_active_account(&account_id, now_timestamp()); + } + })?; + Ok(()) +} + +fn snapshot_existing_managed_account_before_login( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> std::io::Result<()> { + let Some(snapshot) = + persist_current_managed_account_snapshot(codex_home, auth_credentials_store_mode)? + else { + return Ok(()); + }; + upsert_managed_account_snapshot(codex_home, &snapshot, /*set_active*/ false) +} + +fn register_current_managed_account_after_login( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, + managed_account_alias: Option<&str>, +) -> std::io::Result<()> { + let Some(auth) = load_auth_dot_json(codex_home, auth_credentials_store_mode)? else { + return Ok(()); + }; + let Some(snapshot) = + persist_managed_account_auth_snapshot(codex_home, &auth, managed_account_alias)? + else { + return Ok(()); + }; + upsert_managed_account_snapshot(codex_home, &snapshot, /*set_active*/ true) +} + pub async fn login_with_chatgpt( codex_home: PathBuf, forced_chatgpt_workspace_id: Option, cli_auth_credentials_store_mode: AuthCredentialsStoreMode, + managed_account_alias: Option, ) -> std::io::Result<()> { + snapshot_existing_managed_account_before_login(&codex_home, cli_auth_credentials_store_mode)?; + let opts = ServerOptions::new( - codex_home, + codex_home.clone(), CLIENT_ID.to_string(), forced_chatgpt_workspace_id, cli_auth_credentials_store_mode, @@ -125,10 +187,18 @@ pub async fn login_with_chatgpt( print_login_server_start(server.actual_port, &server.auth_url); - server.block_until_done().await + server.block_until_done().await?; + register_current_managed_account_after_login( + &codex_home, + cli_auth_credentials_store_mode, + managed_account_alias.as_deref(), + ) } -pub async fn run_login_with_chatgpt(cli_config_overrides: CliConfigOverrides) -> ! { +pub async fn run_login_with_chatgpt( + cli_config_overrides: CliConfigOverrides, + managed_account_alias: Option, +) -> ! { let config = load_config_or_exit(cli_config_overrides).await; let _login_log_guard = init_login_file_logging(&config); tracing::info!("starting browser login flow"); @@ -144,6 +214,7 @@ pub async fn run_login_with_chatgpt(cli_config_overrides: CliConfigOverrides) -> config.codex_home, forced_chatgpt_workspace_id, config.cli_auth_credentials_store_mode, + managed_account_alias, ) .await { @@ -219,6 +290,7 @@ pub async fn run_login_with_device_code( cli_config_overrides: CliConfigOverrides, issuer_base_url: Option, client_id: Option, + managed_account_alias: Option, ) -> ! { let config = load_config_or_exit(cli_config_overrides).await; let _login_log_guard = init_login_file_logging(&config); @@ -227,6 +299,15 @@ pub async fn run_login_with_device_code( eprintln!("{CHATGPT_LOGIN_DISABLED_MESSAGE}"); std::process::exit(1); } + if let Err(err) = snapshot_existing_managed_account_before_login( + &config.codex_home, + config.cli_auth_credentials_store_mode, + ) { + eprintln!("Error preparing managed account snapshots: {err}"); + std::process::exit(1); + } + + let codex_home = config.codex_home.clone(); let forced_chatgpt_workspace_id = config.forced_chatgpt_workspace_id.clone(); let mut opts = ServerOptions::new( config.codex_home, @@ -239,6 +320,14 @@ pub async fn run_login_with_device_code( } match run_device_code_login(opts).await { Ok(()) => { + if let Err(err) = register_current_managed_account_after_login( + &codex_home, + config.cli_auth_credentials_store_mode, + managed_account_alias.as_deref(), + ) { + eprintln!("Error finalizing managed account login: {err}"); + std::process::exit(1); + } eprintln!("{LOGIN_SUCCESS_MESSAGE}"); std::process::exit(0); } @@ -257,6 +346,7 @@ pub async fn run_login_with_device_code_fallback_to_browser( cli_config_overrides: CliConfigOverrides, issuer_base_url: Option, client_id: Option, + managed_account_alias: Option, ) -> ! { let config = load_config_or_exit(cli_config_overrides).await; let _login_log_guard = init_login_file_logging(&config); @@ -265,7 +355,15 @@ pub async fn run_login_with_device_code_fallback_to_browser( eprintln!("{CHATGPT_LOGIN_DISABLED_MESSAGE}"); std::process::exit(1); } + if let Err(err) = snapshot_existing_managed_account_before_login( + &config.codex_home, + config.cli_auth_credentials_store_mode, + ) { + eprintln!("Error preparing managed account snapshots: {err}"); + std::process::exit(1); + } + let codex_home = config.codex_home.clone(); let forced_chatgpt_workspace_id = config.forced_chatgpt_workspace_id.clone(); let mut opts = ServerOptions::new( config.codex_home, @@ -280,6 +378,14 @@ pub async fn run_login_with_device_code_fallback_to_browser( match run_device_code_login(opts.clone()).await { Ok(()) => { + if let Err(err) = register_current_managed_account_after_login( + &codex_home, + config.cli_auth_credentials_store_mode, + managed_account_alias.as_deref(), + ) { + eprintln!("Error finalizing managed account login: {err}"); + std::process::exit(1); + } eprintln!("{LOGIN_SUCCESS_MESSAGE}"); std::process::exit(0); } @@ -291,6 +397,14 @@ pub async fn run_login_with_device_code_fallback_to_browser( print_login_server_start(server.actual_port, &server.auth_url); match server.block_until_done().await { Ok(()) => { + if let Err(err) = register_current_managed_account_after_login( + &codex_home, + config.cli_auth_credentials_store_mode, + managed_account_alias.as_deref(), + ) { + eprintln!("Error finalizing managed account login: {err}"); + std::process::exit(1); + } eprintln!("{LOGIN_SUCCESS_MESSAGE}"); std::process::exit(0); } @@ -392,6 +506,20 @@ fn safe_format_key(key: &str) -> String { #[cfg(test)] mod tests { + use codex_app_server_protocol::AuthMode as ApiAuthMode; + use codex_core::auth::AuthCredentialsStoreMode; + use codex_core::auth::AuthDotJson; + use codex_core::auth::save_auth; + use codex_core::token_data::TokenData; + use codex_ext::AccountPoolStore; + use codex_ext::ManagedAccountAuthStore; + use pretty_assertions::assert_eq; + use serde_json::json; + use tempfile::tempdir; + + use crate::login::register_current_managed_account_after_login; + use crate::login::snapshot_existing_managed_account_before_login; + use super::safe_format_key; #[test] @@ -405,4 +533,86 @@ mod tests { let key = "sk-proj-12345"; assert_eq!(safe_format_key(key), "***"); } + + fn fake_jwt(email: &str, account_id: &str, plan_type: &str) -> String { + let header = json!({"alg":"none","typ":"JWT"}); + let payload = json!({ + "email": email, + "https://api.openai.com/auth": { + "chatgpt_account_id": account_id, + "chatgpt_plan_type": plan_type, + }, + }); + let encode = |value: serde_json::Value| -> String { + use base64::Engine; + base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(serde_json::to_vec(&value).expect("serialize")) + }; + format!("{}.{}.sig", encode(header), encode(payload)) + } + + fn chatgpt_auth(account_id: &str, email: &str) -> AuthDotJson { + AuthDotJson { + auth_mode: Some(ApiAuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + id_token: codex_core::token_data::parse_chatgpt_jwt_claims(&fake_jwt( + email, account_id, "pro", + )) + .expect("id token"), + access_token: fake_jwt(email, account_id, "pro"), + refresh_token: "refresh-token".to_string(), + account_id: Some(account_id.to_string()), + }), + last_refresh: None, + } + } + + #[test] + fn snapshot_existing_managed_account_before_login_copies_root_auth() { + let tempdir = tempdir().expect("tempdir"); + let auth = chatgpt_auth("acct-existing", "existing@example.com"); + save_auth(tempdir.path(), &auth, AuthCredentialsStoreMode::File).expect("save root auth"); + + snapshot_existing_managed_account_before_login( + tempdir.path(), + AuthCredentialsStoreMode::File, + ) + .expect("snapshot existing auth"); + + let managed_auth = ManagedAccountAuthStore::new(tempdir.path().to_path_buf()) + .load_account_auth("acct-existing") + .expect("managed auth"); + let pool = AccountPoolStore::new(tempdir.path().to_path_buf()) + .load() + .expect("load account pool"); + + assert_eq!(managed_auth, auth); + assert_eq!(pool.accounts[0].alias, "acct-existing"); + } + + #[test] + fn register_current_managed_account_after_login_saves_alias_and_marks_active() { + let tempdir = tempdir().expect("tempdir"); + let auth = chatgpt_auth("acct-primary", "primary@example.com"); + save_auth(tempdir.path(), &auth, AuthCredentialsStoreMode::File).expect("save root auth"); + + register_current_managed_account_after_login( + tempdir.path(), + AuthCredentialsStoreMode::File, + Some("Primary"), + ) + .expect("register managed account"); + + let managed_auth = ManagedAccountAuthStore::new(tempdir.path().to_path_buf()) + .load_account_auth("acct-primary") + .expect("managed auth"); + let pool = AccountPoolStore::new(tempdir.path().to_path_buf()) + .load() + .expect("load account pool"); + + assert_eq!(managed_auth, auth); + assert_eq!(pool.active_account_id.as_deref(), Some("acct-primary")); + assert_eq!(pool.accounts[0].alias, "Primary"); + } } diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 8446e457b..f090309ca 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -275,6 +275,13 @@ struct LoginCommand { #[clap(skip)] config_overrides: CliConfigOverrides, + #[arg( + long = "auth", + value_name = "ALIAS", + help = "Save this ChatGPT login into the managed account pool with the given alias" + )] + auth: Option, + #[arg( long = "with-api-key", help = "Read the API key from stdin (e.g. `printenv OPENAI_API_KEY | codex login --with-api-key`)" @@ -749,6 +756,7 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { login_cli.config_overrides, login_cli.issuer_base_url, login_cli.client_id, + login_cli.auth, ) .await; } else if login_cli.api_key.is_some() { @@ -757,10 +765,14 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { ); std::process::exit(1); } else if login_cli.with_api_key { + if login_cli.auth.is_some() { + eprintln!("The --auth flag is only supported for ChatGPT login flows."); + std::process::exit(1); + } let api_key = read_api_key_from_stdin(); run_login_with_api_key(login_cli.config_overrides, api_key).await; } else { - run_login_with_chatgpt(login_cli.config_overrides).await; + run_login_with_chatgpt(login_cli.config_overrides, login_cli.auth).await; } } } diff --git a/codex-rs/ext/Cargo.toml b/codex-rs/ext/Cargo.toml new file mode 100644 index 000000000..84f3fc634 --- /dev/null +++ b/codex-rs/ext/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "codex-ext" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lib] +name = "codex_ext" +path = "src/lib.rs" + +[lints] +workspace = true + +[dependencies] +codex-core = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } + +[dev-dependencies] +base64 = { workspace = true } +codex-app-server-protocol = { workspace = true } +pretty_assertions = { workspace = true } +tempfile = { workspace = true } diff --git a/codex-rs/ext/README.md b/codex-rs/ext/README.md new file mode 100644 index 000000000..ad0ac8f99 --- /dev/null +++ b/codex-rs/ext/README.md @@ -0,0 +1,14 @@ +# codex-ext + +`codex-ext` is the fork-owned extension layer for long-lived customization. + +Current MVP scope: + +- stable capability-negotiation shapes for future plugin runtimes +- persisted multi-account pool state +- default account router model for threshold-based fallback +- data structures intentionally decoupled from current TUI/core internals + +This crate does not yet load WASM modules. The immediate goal is to keep the +fork-specific policy surface isolated so future upgrades only touch a narrow set +of host integration points. diff --git a/codex-rs/ext/src/account_pool.rs b/codex-rs/ext/src/account_pool.rs new file mode 100644 index 000000000..dd3f87647 --- /dev/null +++ b/codex-rs/ext/src/account_pool.rs @@ -0,0 +1,504 @@ +use serde::Deserialize; +use serde::Serialize; +use std::fs; +use std::io; +use std::path::PathBuf; + +use crate::account_signal::AccountLimitSignal; +use crate::account_signal::AccountRateLimitSnapshot; +pub const ACCOUNT_POOL_STATE_RELATIVE_PATH: &str = "accounts/account-pool.json"; +const ACCOUNT_POOL_STATE_VERSION: u32 = 1; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AccountUsageWindowKind { + FiveHour, + Weekly, + Custom, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum UsageEstimateSource { + Manual, + ResponseErrorInference, + LocalHeuristic, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountUsageWindow { + pub kind: AccountUsageWindowKind, + pub label: String, + pub estimated_used_units: u32, + pub estimated_limit_units: Option, + pub reset_at: Option, + pub source: UsageEstimateSource, +} + +impl AccountUsageWindow { + pub fn pressure_permille(&self) -> Option { + let limit = self.estimated_limit_units?; + if limit == 0 { + return None; + } + + let used = self.estimated_used_units.min(limit); + let permille = (u64::from(used) * 1000) / u64::from(limit); + u16::try_from(permille).ok() + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountRecord { + pub id: String, + pub alias: String, + pub masked_email: Option, + pub plan_label: Option, + pub priority: u32, + pub enabled: bool, + pub cooldown_until: Option, + pub last_limit_error_at: Option, + pub last_selected_at: Option, + pub usage_windows: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountManagementProfile { + pub id: String, + pub alias: Option, + pub masked_email: Option, + pub plan_label: Option, + pub priority: Option, +} + +impl AccountRecord { + pub fn display_name(&self) -> &str { + if self.alias.trim().is_empty() { + &self.id + } else { + &self.alias + } + } + + pub fn is_available_at(&self, now_ts: i64) -> bool { + if !self.enabled { + return false; + } + + match self.cooldown_until { + Some(cooldown_until) => cooldown_until <= now_ts, + None => true, + } + } + + pub fn highest_pressure_permille(&self) -> Option { + self.usage_windows + .iter() + .filter_map(AccountUsageWindow::pressure_permille) + .max() + } + + pub fn usage_summary(&self) -> Option { + let windows: Vec = self + .usage_windows + .iter() + .map(|window| { + let prefix = match window.kind { + AccountUsageWindowKind::FiveHour => "5h", + AccountUsageWindowKind::Weekly => "week", + AccountUsageWindowKind::Custom => window.label.as_str(), + }; + match window.estimated_limit_units { + Some(limit) => format!("{prefix} {}/{}", window.estimated_used_units, limit), + None => format!("{prefix} {}", window.estimated_used_units), + } + }) + .collect(); + if windows.is_empty() { + None + } else { + Some(windows.join(" · ")) + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountPoolState { + pub version: u32, + pub active_account_id: Option, + pub accounts: Vec, +} + +impl Default for AccountPoolState { + fn default() -> Self { + Self { + version: ACCOUNT_POOL_STATE_VERSION, + active_account_id: None, + accounts: Vec::new(), + } + } +} + +impl AccountPoolState { + pub fn upsert_account(&mut self, profile: AccountManagementProfile) -> bool { + if let Some(account) = self + .accounts + .iter_mut() + .find(|account| account.id == profile.id) + { + let mut changed = false; + let next_alias = profile.alias.unwrap_or_else(|| account.alias.clone()); + if account.alias != next_alias { + account.alias = next_alias; + changed = true; + } + if account.masked_email != profile.masked_email { + account.masked_email = profile.masked_email; + changed = true; + } + if account.plan_label != profile.plan_label { + account.plan_label = profile.plan_label; + changed = true; + } + if let Some(priority) = profile.priority + && account.priority != priority + { + account.priority = priority; + changed = true; + } + return changed; + } + + self.accounts.push(AccountRecord { + id: profile.id.clone(), + alias: profile.alias.unwrap_or(profile.id.clone()), + masked_email: profile.masked_email, + plan_label: profile.plan_label, + priority: profile.priority.unwrap_or_else(|| { + u32::try_from(self.accounts.len()).unwrap_or(u32::MAX.saturating_sub(1)) + }), + enabled: true, + cooldown_until: None, + last_limit_error_at: None, + last_selected_at: None, + usage_windows: Vec::new(), + }); + if self.active_account_id.is_none() { + self.active_account_id = Some(profile.id); + } + true + } + + pub fn set_active_account(&mut self, account_id: &str, now_ts: i64) -> bool { + let Some(account) = self + .accounts + .iter_mut() + .find(|account| account.id == account_id) + else { + return false; + }; + account.last_selected_at = Some(now_ts); + let should_change = self.active_account_id.as_deref() != Some(account_id); + self.active_account_id = Some(account_id.to_string()); + should_change || account.last_selected_at == Some(now_ts) + } + + pub fn rename_account_alias(&mut self, account_id: &str, alias: String) -> bool { + let normalized = alias.trim(); + let Some(account) = self + .accounts + .iter_mut() + .find(|account| account.id == account_id) + else { + return false; + }; + let next_alias = if normalized.is_empty() { + account.id.clone() + } else { + normalized.to_string() + }; + if account.alias == next_alias { + false + } else { + account.alias = next_alias; + true + } + } + + pub fn apply_rate_limit_snapshot( + &mut self, + account_id: &str, + snapshot: &AccountRateLimitSnapshot, + ) -> bool { + let Some(account) = self + .accounts + .iter_mut() + .find(|account| account.id == account_id) + else { + return false; + }; + let previous = account.usage_windows.clone(); + account.usage_windows = rate_limit_windows(snapshot); + previous != account.usage_windows + } + + pub fn apply_limit_signal(&mut self, account_id: &str, signal: &AccountLimitSignal) -> bool { + let Some(account) = self + .accounts + .iter_mut() + .find(|account| account.id == account_id) + else { + return false; + }; + let mut changed = false; + if account.last_limit_error_at != Some(signal.recorded_at) { + account.last_limit_error_at = Some(signal.recorded_at); + changed = true; + } + + if signal.cooldown_until.is_some_and(|cooldown_until| { + account + .cooldown_until + .is_none_or(|current| cooldown_until > current) + }) { + account.cooldown_until = signal.cooldown_until; + changed = true; + } + + changed + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountPoolStore { + codex_home: PathBuf, +} + +impl AccountPoolStore { + pub fn new(codex_home: PathBuf) -> Self { + Self { codex_home } + } + + pub fn path(&self) -> PathBuf { + self.codex_home.join(ACCOUNT_POOL_STATE_RELATIVE_PATH) + } + + pub fn load(&self) -> io::Result { + let path = self.path(); + let contents = match fs::read_to_string(&path) { + Ok(contents) => contents, + Err(err) if err.kind() == io::ErrorKind::NotFound => { + return Ok(AccountPoolState::default()); + } + Err(err) => return Err(err), + }; + + serde_json::from_str(&contents).map_err(io::Error::other) + } + + pub fn save(&self, state: &AccountPoolState) -> io::Result<()> { + let path = self.path(); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + + let contents = serde_json::to_string_pretty(state).map_err(io::Error::other)?; + fs::write(path, contents) + } + + pub fn update(&self, updater: F) -> io::Result + where + F: FnOnce(&mut AccountPoolState), + { + let mut state = self.load()?; + updater(&mut state); + self.save(&state)?; + Ok(state) + } +} + +fn rate_limit_windows(snapshot: &AccountRateLimitSnapshot) -> Vec { + [snapshot.primary.as_ref(), snapshot.secondary.as_ref()] + .into_iter() + .flatten() + .map(rate_limit_window_to_usage_window) + .collect() +} + +fn rate_limit_window_to_usage_window( + window: &crate::account_signal::AccountRateLimitWindow, +) -> AccountUsageWindow { + let used_percent = window.used_percent.clamp(0.0, 100.0).round() as u32; + let (kind, label) = match window.window_minutes { + Some(300) => (AccountUsageWindowKind::FiveHour, "5h".to_string()), + Some(10_080) => (AccountUsageWindowKind::Weekly, "week".to_string()), + Some(minutes) => (AccountUsageWindowKind::Custom, format!("{minutes}m")), + None => (AccountUsageWindowKind::Custom, "window".to_string()), + }; + + AccountUsageWindow { + kind, + label, + estimated_used_units: used_percent, + estimated_limit_units: Some(100), + reset_at: window.resets_at, + source: UsageEstimateSource::ResponseErrorInference, + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + use tempfile::tempdir; + + use super::AccountManagementProfile; + use super::AccountPoolState; + use super::AccountPoolStore; + use super::AccountRecord; + use super::AccountUsageWindow; + use super::AccountUsageWindowKind; + use super::UsageEstimateSource; + use crate::account_signal::AccountLimitSignal; + use crate::account_signal::AccountRateLimitSnapshot; + use crate::account_signal::AccountRateLimitWindow; + use crate::account_signal::LimitSignalKind; + + #[test] + fn missing_account_pool_file_returns_default_state() { + let tempdir = tempdir().expect("tempdir"); + let store = AccountPoolStore::new(tempdir.path().to_path_buf()); + + assert_eq!(store.load().expect("load"), AccountPoolState::default()); + } + + #[test] + fn save_round_trips_account_pool_state() { + let tempdir = tempdir().expect("tempdir"); + let store = AccountPoolStore::new(tempdir.path().to_path_buf()); + let state = AccountPoolState { + version: 1, + active_account_id: Some("acc-primary".to_string()), + accounts: vec![AccountRecord { + id: "acc-primary".to_string(), + alias: "Primary".to_string(), + masked_email: Some("pri***@example.com".to_string()), + plan_label: Some("pro".to_string()), + priority: 0, + enabled: true, + cooldown_until: None, + last_limit_error_at: None, + last_selected_at: Some(10), + usage_windows: vec![AccountUsageWindow { + kind: AccountUsageWindowKind::FiveHour, + label: "5h".to_string(), + estimated_used_units: 12, + estimated_limit_units: Some(40), + reset_at: Some(100), + source: UsageEstimateSource::LocalHeuristic, + }], + }], + }; + + store.save(&state).expect("save"); + + assert_eq!(store.load().expect("reload"), state); + } + + #[test] + fn upsert_and_activate_account() { + let mut state = AccountPoolState::default(); + + assert!(state.upsert_account(AccountManagementProfile { + id: "acc-primary".to_string(), + alias: Some("Primary".to_string()), + masked_email: Some("pri***@example.com".to_string()), + plan_label: Some("pro".to_string()), + priority: Some(0), + })); + assert!(state.set_active_account("acc-primary", 123)); + + assert_eq!(state.active_account_id, Some("acc-primary".to_string())); + assert_eq!(state.accounts[0].alias, "Primary"); + assert_eq!(state.accounts[0].last_selected_at, Some(123)); + } + + #[test] + fn applying_rate_limit_snapshot_replaces_usage_windows() { + let mut state = AccountPoolState { + version: 1, + active_account_id: Some("acc-primary".to_string()), + accounts: vec![AccountRecord { + id: "acc-primary".to_string(), + alias: "Primary".to_string(), + masked_email: None, + plan_label: None, + priority: 0, + enabled: true, + cooldown_until: None, + last_limit_error_at: None, + last_selected_at: None, + usage_windows: Vec::new(), + }], + }; + + assert!(state.apply_rate_limit_snapshot( + "acc-primary", + &AccountRateLimitSnapshot { + limit_name: Some("codex".to_string()), + primary: Some(AccountRateLimitWindow { + used_percent: 87.0, + window_minutes: Some(300), + resets_at: Some(500), + }), + secondary: None, + } + )); + + assert_eq!( + state.accounts[0].usage_windows, + vec![AccountUsageWindow { + kind: AccountUsageWindowKind::FiveHour, + label: "5h".to_string(), + estimated_used_units: 87, + estimated_limit_units: Some(100), + reset_at: Some(500), + source: UsageEstimateSource::ResponseErrorInference, + }] + ); + } + + #[test] + fn limit_signal_updates_cooldown() { + let mut state = AccountPoolState { + version: 1, + active_account_id: Some("acc-primary".to_string()), + accounts: vec![AccountRecord { + id: "acc-primary".to_string(), + alias: "Primary".to_string(), + masked_email: None, + plan_label: None, + priority: 0, + enabled: true, + cooldown_until: None, + last_limit_error_at: None, + last_selected_at: None, + usage_windows: Vec::new(), + }], + }; + + assert!(state.apply_limit_signal( + "acc-primary", + &AccountLimitSignal { + kind: LimitSignalKind::UsageLimit, + recorded_at: 100, + cooldown_until: Some(500), + } + )); + + assert_eq!(state.accounts[0].last_limit_error_at, Some(100)); + assert_eq!(state.accounts[0].cooldown_until, Some(500)); + } +} diff --git a/codex-rs/ext/src/account_signal.rs b/codex-rs/ext/src/account_signal.rs new file mode 100644 index 000000000..a62a09a3e --- /dev/null +++ b/codex-rs/ext/src/account_signal.rs @@ -0,0 +1,136 @@ +use serde::Deserialize; +use serde::Serialize; + +const DEFAULT_GENERIC_LIMIT_COOLDOWN_SECS: i64 = 15 * 60; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountRateLimitWindow { + pub used_percent: f64, + pub window_minutes: Option, + pub resets_at: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountRateLimitSnapshot { + pub limit_name: Option, + pub primary: Option, + pub secondary: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LimitSignalKind { + UsageLimit, + RateLimit, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountLimitSignal { + pub kind: LimitSignalKind, + pub recorded_at: i64, + pub cooldown_until: Option, +} + +pub fn infer_limit_signal( + kind: LimitSignalKind, + recorded_at: i64, + snapshot: Option<&AccountRateLimitSnapshot>, +) -> AccountLimitSignal { + let cooldown_until = snapshot + .and_then(|snapshot| blocking_resets_at(snapshot, recorded_at)) + .or_else(|| match kind { + LimitSignalKind::UsageLimit => None, + LimitSignalKind::RateLimit => Some(recorded_at + DEFAULT_GENERIC_LIMIT_COOLDOWN_SECS), + }); + + AccountLimitSignal { + kind, + recorded_at, + cooldown_until, + } +} + +fn blocking_resets_at(snapshot: &AccountRateLimitSnapshot, recorded_at: i64) -> Option { + let mut saturated = Vec::new(); + let mut any_future = Vec::new(); + + for window in [snapshot.primary.as_ref(), snapshot.secondary.as_ref()] + .into_iter() + .flatten() + { + let Some(resets_at) = window + .resets_at + .filter(|resets_at| *resets_at > recorded_at) + else { + continue; + }; + + any_future.push(resets_at); + if window.used_percent >= 99.5 { + saturated.push(resets_at); + } + } + + saturated + .into_iter() + .min() + .or_else(|| any_future.into_iter().min()) +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::AccountLimitSignal; + use super::AccountRateLimitSnapshot; + use super::AccountRateLimitWindow; + use super::LimitSignalKind; + use super::infer_limit_signal; + + #[test] + fn usage_limit_prefers_saturated_window_reset() { + let signal = infer_limit_signal( + LimitSignalKind::UsageLimit, + 100, + Some(&AccountRateLimitSnapshot { + limit_name: Some("codex".to_string()), + primary: Some(AccountRateLimitWindow { + used_percent: 92.0, + window_minutes: Some(300), + resets_at: Some(500), + }), + secondary: Some(AccountRateLimitWindow { + used_percent: 100.0, + window_minutes: Some(10080), + resets_at: Some(900), + }), + }), + ); + + assert_eq!( + signal, + AccountLimitSignal { + kind: LimitSignalKind::UsageLimit, + recorded_at: 100, + cooldown_until: Some(900), + } + ); + } + + #[test] + fn rate_limit_falls_back_to_short_cooldown_without_snapshot() { + let signal = infer_limit_signal(LimitSignalKind::RateLimit, 100, None); + + assert_eq!( + signal, + AccountLimitSignal { + kind: LimitSignalKind::RateLimit, + recorded_at: 100, + cooldown_until: Some(1000), + } + ); + } +} diff --git a/codex-rs/ext/src/host_api.rs b/codex-rs/ext/src/host_api.rs new file mode 100644 index 000000000..de45f2449 --- /dev/null +++ b/codex-rs/ext/src/host_api.rs @@ -0,0 +1,127 @@ +use serde::Deserialize; +use serde::Serialize; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct HostCapability { + pub name: String, + pub version: u32, +} + +impl HostCapability { + pub fn new(name: impl Into, version: u32) -> Self { + Self { + name: name.into(), + version, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CapabilityRequirement { + pub name: String, + pub minimum_version: u32, +} + +impl CapabilityRequirement { + pub fn new(name: impl Into, minimum_version: u32) -> Self { + Self { + name: name.into(), + minimum_version, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct HostCapabilities { + pub capabilities: Vec, +} + +impl Default for HostCapabilities { + fn default() -> Self { + Self::codex_mvp() + } +} + +impl HostCapabilities { + pub fn codex_mvp() -> Self { + Self { + capabilities: vec![ + HostCapability::new("app-start", /*version*/ 1), + HostCapability::new("session-start", /*version*/ 1), + HostCapability::new("before-turn-start", /*version*/ 1), + HostCapability::new("before-tool-call", /*version*/ 1), + HostCapability::new("after-tool-call", /*version*/ 1), + HostCapability::new("account-routing", /*version*/ 1), + HostCapability::new("control-panel", /*version*/ 1), + ], + } + } + + pub fn supports(&self, requirement: &CapabilityRequirement) -> bool { + self.capabilities.iter().any(|capability| { + capability.name == requirement.name && capability.version >= requirement.minimum_version + }) + } + + pub fn negotiate(&self, requirements: &[CapabilityRequirement]) -> PluginNegotiation { + let mut accepted = Vec::new(); + let mut missing = Vec::new(); + + for requirement in requirements { + if self.supports(requirement) { + accepted.push(requirement.clone()); + } else { + missing.push(requirement.clone()); + } + } + + PluginNegotiation { accepted, missing } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PluginNegotiation { + pub accepted: Vec, + pub missing: Vec, +} + +impl PluginNegotiation { + pub fn is_compatible(&self) -> bool { + self.missing.is_empty() + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::CapabilityRequirement; + use super::HostCapabilities; + + #[test] + fn codex_mvp_capabilities_support_account_routing() { + let capabilities = HostCapabilities::codex_mvp(); + let negotiation = capabilities.negotiate(&[ + CapabilityRequirement::new("account-routing", 1), + CapabilityRequirement::new("control-panel", 1), + ]); + + assert!(negotiation.is_compatible()); + assert_eq!(negotiation.accepted.len(), 2); + assert_eq!(negotiation.missing.len(), 0); + } + + #[test] + fn negotiation_reports_missing_versions() { + let capabilities = HostCapabilities::codex_mvp(); + let negotiation = + capabilities.negotiate(&[CapabilityRequirement::new("before-turn-start", 2)]); + + assert!(!negotiation.is_compatible()); + assert_eq!(negotiation.accepted.len(), 0); + assert_eq!( + negotiation.missing, + vec![CapabilityRequirement::new("before-turn-start", 2)] + ); + } +} diff --git a/codex-rs/ext/src/lib.rs b/codex-rs/ext/src/lib.rs new file mode 100644 index 000000000..c269ee4eb --- /dev/null +++ b/codex-rs/ext/src/lib.rs @@ -0,0 +1,34 @@ +pub mod account_pool; +pub mod account_signal; +pub mod host_api; +pub mod managed_account_auth; +pub mod router; + +pub use account_pool::AccountManagementProfile; +pub use account_pool::AccountPoolState; +pub use account_pool::AccountPoolStore; +pub use account_pool::AccountRecord; +pub use account_pool::AccountUsageWindow; +pub use account_pool::AccountUsageWindowKind; +pub use account_pool::UsageEstimateSource; +pub use account_signal::AccountLimitSignal; +pub use account_signal::AccountRateLimitSnapshot; +pub use account_signal::AccountRateLimitWindow; +pub use account_signal::LimitSignalKind; +pub use account_signal::infer_limit_signal; +pub use host_api::CapabilityRequirement; +pub use host_api::HostCapabilities; +pub use host_api::HostCapability; +pub use host_api::PluginNegotiation; +pub use managed_account_auth::MANAGED_ACCOUNTS_RELATIVE_DIR; +pub use managed_account_auth::ManagedAccountAuthStore; +pub use managed_account_auth::ManagedAccountSnapshot; +pub use managed_account_auth::activate_managed_account; +pub use managed_account_auth::load_current_managed_account_snapshot; +pub use managed_account_auth::persist_current_managed_account_snapshot; +pub use managed_account_auth::persist_managed_account_auth_snapshot; +pub use router::AccountRouterDecision; +pub use router::AccountRouterDecisionReason; +pub use router::DefaultAccountRouter; +pub use router::RouteTurnRequest; +pub use router::RoutingTrigger; diff --git a/codex-rs/ext/src/managed_account_auth.rs b/codex-rs/ext/src/managed_account_auth.rs new file mode 100644 index 000000000..6de0bf3fb --- /dev/null +++ b/codex-rs/ext/src/managed_account_auth.rs @@ -0,0 +1,284 @@ +use codex_core::auth::AuthCredentialsStoreMode; +use codex_core::auth::AuthDotJson; +use codex_core::auth::load_auth_dot_json; +use codex_core::auth::save_auth; +use std::fs; +use std::fs::OpenOptions; +use std::io; +use std::io::Write; +#[cfg(unix)] +use std::os::unix::fs::OpenOptionsExt; +use std::path::Path; +use std::path::PathBuf; + +use crate::account_pool::AccountManagementProfile; + +pub const MANAGED_ACCOUNTS_RELATIVE_DIR: &str = "accounts"; +const ACCOUNT_AUTH_FILE_NAME: &str = "auth.json"; + +#[derive(Debug, Clone, PartialEq)] +pub struct ManagedAccountSnapshot { + pub profile: AccountManagementProfile, + pub auth: AuthDotJson, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ManagedAccountAuthStore { + codex_home: PathBuf, +} + +impl ManagedAccountAuthStore { + pub fn new(codex_home: PathBuf) -> Self { + Self { codex_home } + } + + pub fn directory(&self) -> PathBuf { + self.codex_home.join(MANAGED_ACCOUNTS_RELATIVE_DIR) + } + + pub fn account_dir(&self, account_id: &str) -> PathBuf { + self.directory().join(account_id) + } + + pub fn account_auth_path(&self, account_id: &str) -> PathBuf { + self.account_dir(account_id).join(ACCOUNT_AUTH_FILE_NAME) + } + + pub fn load_account_auth(&self, account_id: &str) -> io::Result { + let path = self.account_auth_path(account_id); + let contents = fs::read_to_string(path)?; + serde_json::from_str(&contents).map_err(io::Error::other) + } + + pub fn save_account_auth(&self, account_id: &str, auth: &AuthDotJson) -> io::Result<()> { + let path = self.account_auth_path(account_id); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + + let contents = serde_json::to_string_pretty(auth).map_err(io::Error::other)?; + let mut options = OpenOptions::new(); + options.create(true).truncate(true).write(true); + #[cfg(unix)] + { + options.mode(0o600); + } + let mut file = options.open(path)?; + file.write_all(contents.as_bytes())?; + file.flush() + } +} + +pub fn activate_managed_account( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, + account_id: &str, +) -> io::Result<()> { + let store = ManagedAccountAuthStore::new(codex_home.to_path_buf()); + let auth = store.load_account_auth(account_id)?; + save_auth(codex_home, &auth, auth_credentials_store_mode) +} + +pub fn persist_current_managed_account_snapshot( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> io::Result> { + let Some(snapshot) = + load_current_managed_account_snapshot(codex_home, auth_credentials_store_mode)? + else { + return Ok(None); + }; + let store = ManagedAccountAuthStore::new(codex_home.to_path_buf()); + store.save_account_auth(&snapshot.profile.id, &snapshot.auth)?; + Ok(Some(snapshot)) +} + +pub fn persist_managed_account_auth_snapshot( + codex_home: &Path, + auth: &AuthDotJson, + alias_override: Option<&str>, +) -> io::Result> { + let Some(snapshot) = managed_account_snapshot_from_auth(auth, alias_override) else { + return Ok(None); + }; + let store = ManagedAccountAuthStore::new(codex_home.to_path_buf()); + store.save_account_auth(&snapshot.profile.id, &snapshot.auth)?; + Ok(Some(snapshot)) +} + +pub fn load_current_managed_account_snapshot( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> io::Result> { + let Some(auth_json) = load_auth_dot_json(codex_home, auth_credentials_store_mode)? else { + return Ok(None); + }; + Ok(managed_account_snapshot_from_auth( + &auth_json, /*alias_override*/ None, + )) +} + +fn managed_account_snapshot_from_auth( + auth: &AuthDotJson, + alias_override: Option<&str>, +) -> Option { + let tokens = auth.tokens.as_ref()?; + let id = tokens + .account_id + .clone() + .or_else(|| tokens.id_token.chatgpt_account_id.clone())?; + let alias = alias_override + .map(str::trim) + .filter(|alias| !alias.is_empty()) + .map(ToString::to_string) + .unwrap_or_else(|| id.clone()); + let profile = AccountManagementProfile { + id, + alias: Some(alias), + masked_email: tokens.id_token.email.as_deref().map(mask_email), + plan_label: tokens + .id_token + .get_chatgpt_plan_type() + .map(|plan_type| plan_type.to_ascii_lowercase()), + priority: None, + }; + Some(ManagedAccountSnapshot { + profile, + auth: auth.clone(), + }) +} + +fn mask_email(email: &str) -> String { + let mut parts = email.split('@'); + let local = parts.next().unwrap_or_default(); + let domain = parts.next().unwrap_or_default(); + if local.is_empty() || domain.is_empty() { + return email.to_string(); + } + + let prefix: String = local.chars().take(3).collect(); + format!("{prefix}***@{domain}") +} + +#[cfg(test)] +mod tests { + use super::ManagedAccountAuthStore; + use super::activate_managed_account; + use super::persist_current_managed_account_snapshot; + use super::persist_managed_account_auth_snapshot; + use pretty_assertions::assert_eq; + use tempfile::tempdir; + + use codex_app_server_protocol::AuthMode; + use codex_core::auth::AuthCredentialsStoreMode; + use codex_core::auth::AuthDotJson; + use codex_core::auth::load_auth_dot_json; + use codex_core::auth::save_auth; + use codex_core::token_data::TokenData; + use serde_json::json; + + fn chatgpt_auth(account_id: &str, email: &str) -> AuthDotJson { + AuthDotJson { + auth_mode: Some(AuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + id_token: codex_core::token_data::parse_chatgpt_jwt_claims(&fake_jwt( + email, account_id, "pro", + )) + .expect("id token"), + access_token: fake_jwt(email, account_id, "pro"), + refresh_token: "refresh-token".to_string(), + account_id: Some(account_id.to_string()), + }), + last_refresh: None, + } + } + + fn fake_jwt(email: &str, account_id: &str, plan_type: &str) -> String { + let header = json!({"alg":"none","typ":"JWT"}); + let payload = json!({ + "email": email, + "https://api.openai.com/auth": { + "chatgpt_account_id": account_id, + "chatgpt_plan_type": plan_type, + }, + }); + let encode = |value: serde_json::Value| -> String { + use base64::Engine; + base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(serde_json::to_vec(&value).expect("serialize")) + }; + format!("{}.{}.sig", encode(header), encode(payload)) + } + + #[test] + fn persist_snapshot_writes_account_auth_under_accounts_directory() { + let tempdir = tempdir().expect("tempdir"); + let auth = chatgpt_auth("workspace-1", "user@example.com"); + save_auth(tempdir.path(), &auth, AuthCredentialsStoreMode::File).expect("save auth"); + + let snapshot = persist_current_managed_account_snapshot( + tempdir.path(), + AuthCredentialsStoreMode::File, + ) + .expect("persist snapshot") + .expect("snapshot"); + + assert_eq!(snapshot.profile.id, "workspace-1"); + assert_eq!( + ManagedAccountAuthStore::new(tempdir.path().to_path_buf()) + .load_account_auth("workspace-1") + .expect("load snapshot"), + auth + ); + } + + #[test] + fn persist_managed_account_auth_snapshot_uses_alias_override() { + let tempdir = tempdir().expect("tempdir"); + let auth = chatgpt_auth("workspace-1", "user@example.com"); + + let snapshot = + persist_managed_account_auth_snapshot(tempdir.path(), &auth, Some("Primary")) + .expect("persist snapshot") + .expect("snapshot"); + + assert_eq!(snapshot.profile.alias, Some("Primary".to_string())); + assert_eq!( + ManagedAccountAuthStore::new(tempdir.path().to_path_buf()) + .load_account_auth("workspace-1") + .expect("load snapshot"), + auth + ); + } + + #[test] + fn activate_managed_account_restores_root_auth() { + let tempdir = tempdir().expect("tempdir"); + let primary = chatgpt_auth("workspace-1", "primary@example.com"); + let backup = chatgpt_auth("workspace-2", "backup@example.com"); + let store = ManagedAccountAuthStore::new(tempdir.path().to_path_buf()); + store + .save_account_auth("workspace-1", &primary) + .expect("save primary"); + store + .save_account_auth("workspace-2", &backup) + .expect("save backup"); + + save_auth(tempdir.path(), &primary, AuthCredentialsStoreMode::File) + .expect("save root auth"); + activate_managed_account( + tempdir.path(), + AuthCredentialsStoreMode::File, + "workspace-2", + ) + .expect("activate backup"); + + assert_eq!( + load_auth_dot_json(tempdir.path(), AuthCredentialsStoreMode::File) + .expect("load root auth") + .expect("root auth"), + backup + ); + } +} diff --git a/codex-rs/ext/src/router.rs b/codex-rs/ext/src/router.rs new file mode 100644 index 000000000..7ae93bc4a --- /dev/null +++ b/codex-rs/ext/src/router.rs @@ -0,0 +1,300 @@ +use crate::account_pool::AccountPoolState; +use crate::account_pool::AccountRecord; +use serde::Deserialize; +use serde::Serialize; + +const DEFAULT_USAGE_THRESHOLD_PERMILLE: u16 = 850; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RoutingTrigger { + NormalTurn, + RetryAfterHardError, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AccountRouterDecisionReason { + KeepActiveAccount, + RetryWithFallbackAccount, + ActiveAccountCoolingDown, + ActiveAccountOverThreshold, + PreferredFallbackSelected, + NoHealthyAccount, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RouteTurnRequest { + pub now_ts: i64, + pub trigger: RoutingTrigger, + pub active_account_id: Option, + pub preferred_account_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountRouterDecision { + pub account_id: Option, + pub reason: AccountRouterDecisionReason, + pub retry_immediately: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DefaultAccountRouter { + usage_threshold_permille: u16, +} + +impl Default for DefaultAccountRouter { + fn default() -> Self { + Self { + usage_threshold_permille: DEFAULT_USAGE_THRESHOLD_PERMILLE, + } + } +} + +impl DefaultAccountRouter { + pub fn new(usage_threshold_permille: u16) -> Self { + Self { + usage_threshold_permille, + } + } + + pub fn select_account( + &self, + state: &AccountPoolState, + request: &RouteTurnRequest, + ) -> AccountRouterDecision { + let active = request + .active_account_id + .as_deref() + .or(state.active_account_id.as_deref()) + .and_then(|account_id| { + state + .accounts + .iter() + .find(|account| account.id == account_id) + }); + + if request.trigger == RoutingTrigger::NormalTurn + && let Some(active) = active + && self.is_healthy(active, request.now_ts) + { + return AccountRouterDecision { + account_id: Some(active.id.clone()), + reason: AccountRouterDecisionReason::KeepActiveAccount, + retry_immediately: false, + }; + } + + if request.trigger == RoutingTrigger::RetryAfterHardError { + if let Some(fallback) = self.pick_fallback(state, request, active) { + return AccountRouterDecision { + account_id: Some(fallback.id.clone()), + reason: AccountRouterDecisionReason::RetryWithFallbackAccount, + retry_immediately: true, + }; + } + + return AccountRouterDecision { + account_id: None, + reason: AccountRouterDecisionReason::NoHealthyAccount, + retry_immediately: false, + }; + } + + if let Some(active) = active { + let reason = if active.is_available_at(request.now_ts) { + AccountRouterDecisionReason::ActiveAccountOverThreshold + } else { + AccountRouterDecisionReason::ActiveAccountCoolingDown + }; + + if let Some(fallback) = self.pick_fallback(state, request, Some(active)) { + return AccountRouterDecision { + account_id: Some(fallback.id.clone()), + reason, + retry_immediately: false, + }; + } + } + + if let Some(preferred) = request + .preferred_account_id + .as_deref() + .and_then(|account_id| { + state + .accounts + .iter() + .find(|account| account.id == account_id) + }) + .filter(|account| self.is_healthy(account, request.now_ts)) + { + return AccountRouterDecision { + account_id: Some(preferred.id.clone()), + reason: AccountRouterDecisionReason::PreferredFallbackSelected, + retry_immediately: false, + }; + } + + AccountRouterDecision { + account_id: None, + reason: AccountRouterDecisionReason::NoHealthyAccount, + retry_immediately: false, + } + } + + fn is_healthy(&self, account: &AccountRecord, now_ts: i64) -> bool { + account.is_available_at(now_ts) + && account + .highest_pressure_permille() + .is_none_or(|pressure| pressure < self.usage_threshold_permille) + } + + fn pick_fallback<'a>( + &self, + state: &'a AccountPoolState, + request: &RouteTurnRequest, + active: Option<&AccountRecord>, + ) -> Option<&'a AccountRecord> { + let active_id = active.map(|account| account.id.as_str()); + let mut candidates: Vec<&AccountRecord> = state + .accounts + .iter() + .filter(|account| Some(account.id.as_str()) != active_id) + .filter(|account| self.is_healthy(account, request.now_ts)) + .collect(); + + candidates.sort_by_key(|account| (account.priority, account.last_selected_at.unwrap_or(0))); + candidates.into_iter().next() + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use crate::account_pool::AccountPoolState; + use crate::account_pool::AccountRecord; + use crate::account_pool::AccountUsageWindow; + use crate::account_pool::AccountUsageWindowKind; + use crate::account_pool::UsageEstimateSource; + + use super::AccountRouterDecision; + use super::AccountRouterDecisionReason; + use super::DefaultAccountRouter; + use super::RouteTurnRequest; + use super::RoutingTrigger; + + fn account( + id: &str, + priority: u32, + used: u32, + limit: u32, + cooldown_until: Option, + ) -> AccountRecord { + AccountRecord { + id: id.to_string(), + alias: id.to_string(), + masked_email: None, + plan_label: None, + priority, + enabled: true, + cooldown_until, + last_limit_error_at: None, + last_selected_at: None, + usage_windows: vec![AccountUsageWindow { + kind: AccountUsageWindowKind::FiveHour, + label: "5h".to_string(), + estimated_used_units: used, + estimated_limit_units: Some(limit), + reset_at: None, + source: UsageEstimateSource::LocalHeuristic, + }], + } + } + + #[test] + fn normal_turn_keeps_healthy_active_account() { + let router = DefaultAccountRouter::default(); + let state = AccountPoolState { + version: 1, + active_account_id: Some("primary".to_string()), + accounts: vec![ + account("primary", 0, 10, 40, None), + account("backup", 1, 1, 40, None), + ], + }; + let request = RouteTurnRequest { + now_ts: 100, + trigger: RoutingTrigger::NormalTurn, + active_account_id: Some("primary".to_string()), + preferred_account_id: None, + }; + + assert_eq!( + router.select_account(&state, &request), + AccountRouterDecision { + account_id: Some("primary".to_string()), + reason: AccountRouterDecisionReason::KeepActiveAccount, + retry_immediately: false, + } + ); + } + + #[test] + fn retry_after_hard_error_picks_fallback_account() { + let router = DefaultAccountRouter::default(); + let state = AccountPoolState { + version: 1, + active_account_id: Some("primary".to_string()), + accounts: vec![ + account("primary", 0, 38, 40, Some(200)), + account("backup", 1, 1, 40, None), + ], + }; + let request = RouteTurnRequest { + now_ts: 100, + trigger: RoutingTrigger::RetryAfterHardError, + active_account_id: Some("primary".to_string()), + preferred_account_id: None, + }; + + assert_eq!( + router.select_account(&state, &request), + AccountRouterDecision { + account_id: Some("backup".to_string()), + reason: AccountRouterDecisionReason::RetryWithFallbackAccount, + retry_immediately: true, + } + ); + } + + #[test] + fn normal_turn_switches_when_active_is_over_threshold() { + let router = DefaultAccountRouter::default(); + let state = AccountPoolState { + version: 1, + active_account_id: Some("primary".to_string()), + accounts: vec![ + account("primary", 0, 35, 40, None), + account("backup", 1, 4, 40, None), + ], + }; + let request = RouteTurnRequest { + now_ts: 100, + trigger: RoutingTrigger::NormalTurn, + active_account_id: Some("primary".to_string()), + preferred_account_id: None, + }; + + assert_eq!( + router.select_account(&state, &request), + AccountRouterDecision { + account_id: Some("backup".to_string()), + reason: AccountRouterDecisionReason::ActiveAccountOverThreshold, + retry_immediately: false, + } + ); + } +} diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index b726eeed8..02e673bb9 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -1084,13 +1084,12 @@ pub(crate) async fn obtain_api_key( } #[cfg(test)] mod tests { - use pretty_assertions::assert_eq; - use super::TokenEndpointErrorDetail; use super::parse_token_endpoint_error; use super::redact_sensitive_query_value; use super::redact_sensitive_url_parts; use super::sanitize_url_for_logging; + use pretty_assertions::assert_eq; #[test] fn parse_token_endpoint_error_prefers_error_description() { diff --git a/codex-rs/tui/Cargo.toml b/codex-rs/tui/Cargo.toml index 8013b1325..0457d045d 100644 --- a/codex-rs/tui/Cargo.toml +++ b/codex-rs/tui/Cargo.toml @@ -37,6 +37,7 @@ codex-chatgpt = { workspace = true } codex-client = { workspace = true } codex-cloud-requirements = { workspace = true } codex-core = { workspace = true } +codex-ext = { workspace = true } codex-features = { workspace = true } codex-feedback = { workspace = true } codex-file-search = { workspace = true } diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 4aa51df21..b1f7b3045 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -38,6 +38,7 @@ use crate::tui; use crate::tui::TuiEvent; use crate::update_action::UpdateAction; use crate::version::CODEX_CLI_VERSION; +use chrono::Utc; use codex_ansi_escape::ansi_escape_line; use codex_app_server_client::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY; use codex_app_server_client::InProcessAppServerClient; @@ -70,6 +71,10 @@ use codex_core::models_manager::model_presets::HIDE_GPT_5_1_CODEX_MAX_MIGRATION_ use codex_core::models_manager::model_presets::HIDE_GPT5_1_MIGRATION_PROMPT_CONFIG; #[cfg(target_os = "windows")] use codex_core::windows_sandbox::WindowsSandboxLevelExt; +use codex_ext::AccountPoolStore; +use codex_ext::HostCapabilities; +use codex_ext::activate_managed_account; +use codex_ext::persist_current_managed_account_snapshot; use codex_features::Feature; use codex_otel::SessionTelemetry; use codex_otel::TelemetryAuthMode; @@ -1916,6 +1921,298 @@ impl App { self.sync_active_agent_label(); } + fn open_control_panel(&mut self) { + let capabilities = HostCapabilities::codex_mvp(); + let items = vec![ + SelectionItem { + name: "Sessions".to_string(), + description: Some("Resume or switch saved chats.".to_string()), + actions: vec![Box::new(|tx| tx.send(AppEvent::OpenResumePickerAll))], + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Fork Current Session".to_string(), + description: Some("Fork the current thread into a new session.".to_string()), + actions: vec![Box::new(|tx| tx.send(AppEvent::ForkCurrentSession))], + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Accounts".to_string(), + description: Some("Inspect the managed multi-account pool.".to_string()), + actions: vec![Box::new(|tx| tx.send(AppEvent::OpenAccountsPanel))], + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Undo Last User Message".to_string(), + description: Some( + "Planned MVP follow-up: restore the last sent input and roll back one turn." + .to_string(), + ), + is_disabled: true, + disabled_reason: Some("Not wired yet.".to_string()), + ..Default::default() + }, + ]; + + self.chat_widget.show_selection_view(SelectionViewParams { + view_id: Some("fork-control-panel"), + title: Some("Control Panel".to_string()), + subtitle: Some(format!( + "Fork extension host: {} negotiated MVP capabilities.", + capabilities.capabilities.len() + )), + footer_hint: Some(standard_popup_hint_line()), + items, + ..Default::default() + }); + } + + fn open_accounts_panel(&mut self) { + let store = AccountPoolStore::new(self.config.codex_home.clone()); + let path = store.path(); + let state = self.sync_current_auth_into_account_pool(&store); + let mut items = Vec::new(); + let subtitle = match &state { + Ok(state) => Some(format!( + "{} account(s) configured. Active: {}.", + state.accounts.len(), + state.active_account_id.as_deref().unwrap_or("none") + )), + Err(err) => Some(format!("Failed to read account pool: {err}")), + }; + + match state { + Ok(state) => { + for account in &state.accounts { + let mut description_parts = Vec::new(); + if let Some(masked_email) = &account.masked_email { + description_parts.push(masked_email.clone()); + } + if let Some(plan_label) = &account.plan_label { + description_parts.push(plan_label.clone()); + } + if let Some(usage_summary) = account.usage_summary() { + description_parts.push(usage_summary); + } + if let Some(cooldown_until) = account.cooldown_until { + description_parts.push(format!("cooldown until {cooldown_until}")); + } + + let account_id = account.id.clone(); + let search_value = format!( + "{} {} {}", + account.display_name(), + account.id, + account.masked_email.clone().unwrap_or_default() + ); + items.push(SelectionItem { + name: account.display_name().to_string(), + description: (!description_parts.is_empty()) + .then(|| description_parts.join(" · ")), + is_current: state.active_account_id.as_deref() == Some(account.id.as_str()), + is_disabled: !account.enabled, + disabled_reason: (!account.enabled).then(|| "Disabled".to_string()), + actions: vec![Box::new(move |tx| { + tx.send(AppEvent::SetManagedAccountActive(account_id.clone())); + })], + dismiss_on_select: true, + search_value: Some(search_value), + ..Default::default() + }); + } + if !state.accounts.is_empty() { + items.push(SelectionItem { + name: "Rename".to_string(), + description: Some( + "Open alias rename actions for managed accounts.".to_string(), + ), + actions: vec![Box::new(|tx| { + tx.send(AppEvent::OpenManagedAccountRenamePanel) + })], + dismiss_on_select: true, + ..Default::default() + }); + } + } + Err(err) => { + items.push(SelectionItem { + name: "Account pool unavailable".to_string(), + description: Some(err.to_string()), + is_disabled: true, + ..Default::default() + }); + } + } + + if items.is_empty() { + items.push(SelectionItem { + name: "No managed accounts yet".to_string(), + description: Some(format!( + "Run repeated ChatGPT login flows and persist them into {}.", + path.display() + )), + is_disabled: true, + ..Default::default() + }); + } + + self.chat_widget.show_selection_view(SelectionViewParams { + view_id: Some("fork-accounts-panel"), + title: Some("Accounts".to_string()), + subtitle, + footer_hint: Some(standard_popup_hint_line()), + footer_note: Some(Line::from(path.display().to_string()).dim()), + items, + ..Default::default() + }); + } + + fn open_managed_account_rename_panel(&mut self) { + let store = AccountPoolStore::new(self.config.codex_home.clone()); + let path = store.path(); + let state = self.sync_current_auth_into_account_pool(&store); + let mut items = Vec::new(); + let subtitle = match &state { + Ok(state) => Some(format!( + "Rename aliases for {} managed account(s).", + state.accounts.len() + )), + Err(err) => Some(format!("Failed to read account pool: {err}")), + }; + + match state { + Ok(state) => { + for account in &state.accounts { + let account_id = account.id.clone(); + let current_alias = account.alias.clone(); + let mut description_parts = Vec::new(); + if let Some(alias) = + (!current_alias.is_empty()).then_some(current_alias.clone()) + { + description_parts.push(format!("alias: {alias}")); + } + description_parts.push(account.id.clone()); + if let Some(masked_email) = &account.masked_email { + description_parts.push(masked_email.clone()); + } + items.push(SelectionItem { + name: format!("Rename {}", account.display_name()), + description: Some(description_parts.join(" · ")), + actions: vec![Box::new(move |tx| { + tx.send(AppEvent::OpenRenameManagedAccountAliasPrompt { + account_id: account_id.clone(), + current_alias: current_alias.clone(), + }); + })], + dismiss_on_select: true, + ..Default::default() + }); + } + } + Err(err) => { + items.push(SelectionItem { + name: "Account pool unavailable".to_string(), + description: Some(err.to_string()), + is_disabled: true, + ..Default::default() + }); + } + } + + if items.is_empty() { + items.push(SelectionItem { + name: "No managed accounts yet".to_string(), + description: Some(format!( + "Run repeated ChatGPT login flows and persist them into {}.", + path.display() + )), + is_disabled: true, + ..Default::default() + }); + } + + self.chat_widget.show_selection_view(SelectionViewParams { + view_id: Some("fork-account-rename-panel"), + title: Some("Rename".to_string()), + subtitle, + footer_hint: Some(standard_popup_hint_line()), + footer_note: Some(Line::from(path.display().to_string()).dim()), + items, + ..Default::default() + }); + } + + fn sync_current_auth_into_account_pool( + &self, + store: &AccountPoolStore, + ) -> std::io::Result { + let snapshot = persist_current_managed_account_snapshot( + &self.config.codex_home, + self.config.cli_auth_credentials_store_mode, + )?; + store.update(|state| { + if let Some(snapshot) = &snapshot { + state.upsert_account(snapshot.profile.clone()); + } + }) + } + + fn set_managed_account_active(&mut self, account_id: String) { + if let Err(err) = persist_current_managed_account_snapshot( + &self.config.codex_home, + self.config.cli_auth_credentials_store_mode, + ) { + self.chat_widget + .add_to_history(history_cell::new_error_event(format!( + "Failed to snapshot current managed account auth: {err}" + ))); + return; + } + if let Err(err) = activate_managed_account( + &self.config.codex_home, + self.config.cli_auth_credentials_store_mode, + &account_id, + ) { + self.chat_widget + .add_to_history(history_cell::new_error_event(format!( + "Failed to activate managed account auth: {err}" + ))); + return; + } + self.auth_manager.reload(); + + let store = AccountPoolStore::new(self.config.codex_home.clone()); + if let Err(err) = store.update(|state| { + state.set_active_account(&account_id, Utc::now().timestamp()); + }) { + self.chat_widget + .add_to_history(history_cell::new_error_event(format!( + "Failed to update active managed account: {err}" + ))); + return; + } + + self.open_accounts_panel(); + } + + fn save_managed_account_alias(&mut self, account_id: String, alias: String) { + let store = AccountPoolStore::new(self.config.codex_home.clone()); + if let Err(err) = store.update(|state| { + state.rename_account_alias(&account_id, alias.clone()); + }) { + self.chat_widget + .add_to_history(history_cell::new_error_event(format!( + "Failed to save managed account alias: {err}" + ))); + return; + } + + self.open_managed_account_rename_panel(); + } + /// Marks a cached picker thread closed and recomputes the contextual footer label. /// /// Closing a thread is not the same as removing it: users can still inspect finished agent @@ -2646,6 +2943,28 @@ impl App { async fn handle_event(&mut self, tui: &mut tui::Tui, event: AppEvent) -> Result { match event { + AppEvent::OpenControlPanel => { + self.open_control_panel(); + } + AppEvent::OpenAccountsPanel => { + self.open_accounts_panel(); + } + AppEvent::OpenManagedAccountRenamePanel => { + self.open_managed_account_rename_panel(); + } + AppEvent::SetManagedAccountActive(account_id) => { + self.set_managed_account_active(account_id); + } + AppEvent::OpenRenameManagedAccountAliasPrompt { + account_id, + current_alias, + } => { + self.chat_widget + .open_managed_account_alias_prompt(account_id, current_alias); + } + AppEvent::SaveManagedAccountAlias { account_id, alias } => { + self.save_managed_account_alias(account_id, alias); + } AppEvent::NewSession => { self.start_fresh_session_with_summary_hint(tui).await; } @@ -2754,6 +3073,105 @@ impl App { // Leaving alt-screen may blank the inline viewport; force a redraw either way. tui.frame_requester().schedule_frame(); } + AppEvent::OpenResumePickerAll => { + match crate::resume_picker::run_resume_picker( + tui, + &self.config, + /*show_all*/ true, + ) + .await? + { + SessionSelection::Resume(target_session) => { + let current_cwd = self.config.cwd.clone(); + let resume_cwd = match crate::resolve_cwd_for_resume_or_fork( + tui, + &self.config, + ¤t_cwd, + target_session.thread_id, + &target_session.path, + CwdPromptAction::Resume, + /*allow_prompt*/ true, + ) + .await? + { + crate::ResolveCwdOutcome::Continue(Some(cwd)) => cwd, + crate::ResolveCwdOutcome::Continue(None) => current_cwd.clone(), + crate::ResolveCwdOutcome::Exit => { + return Ok(AppRunControl::Exit(ExitReason::UserRequested)); + } + }; + let mut resume_config = match self + .rebuild_config_for_resume_or_fallback(¤t_cwd, resume_cwd) + .await + { + Ok(cfg) => cfg, + Err(err) => { + self.chat_widget.add_error_message(format!( + "Failed to rebuild configuration for resume: {err}" + )); + return Ok(AppRunControl::Continue); + } + }; + self.apply_runtime_policy_overrides(&mut resume_config); + let summary = session_summary( + self.chat_widget.token_usage(), + self.chat_widget.thread_id(), + self.chat_widget.thread_name(), + ); + match self + .server + .resume_thread_from_rollout( + resume_config.clone(), + target_session.path.clone(), + self.auth_manager.clone(), + /*parent_trace*/ None, + ) + .await + { + Ok(resumed) => { + self.shutdown_current_thread().await; + self.config = resume_config; + tui.set_notification_method(self.config.tui_notification_method); + self.file_search.update_search_dir(self.config.cwd.clone()); + let init = self.chatwidget_init_for_forked_or_resumed_thread( + tui, + self.config.clone(), + ); + self.replace_chat_widget(ChatWidget::new_from_existing( + init, + resumed.thread, + resumed.session_configured, + )); + self.reset_thread_event_state(); + if let Some(summary) = summary { + let mut lines: Vec> = + vec![summary.usage_line.clone().into()]; + if let Some(command) = summary.resume_command { + let spans = vec![ + "To continue this session, run ".into(), + command.cyan(), + ]; + lines.push(spans.into()); + } + self.chat_widget.add_plain_history_lines(lines); + } + } + Err(err) => { + let path_display = target_session.path.display(); + self.chat_widget.add_error_message(format!( + "Failed to resume session from {path_display}: {err}" + )); + } + } + } + SessionSelection::Exit + | SessionSelection::StartFresh + | SessionSelection::Fork(_) => {} + } + + // Leaving alt-screen may blank the inline viewport; force a redraw either way. + tui.frame_requester().schedule_frame(); + } AppEvent::ForkCurrentSession => { self.session_telemetry.counter( "codex.thread.fork", diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index 71fc7be27..8cc167837 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -74,6 +74,24 @@ pub(crate) struct ConnectorsSnapshot { #[derive(Debug)] pub(crate) enum AppEvent { CodexEvent(Event), + /// Open the fork-owned control panel. + OpenControlPanel, + /// Open the account pool panel inside the control panel flow. + OpenAccountsPanel, + /// Open the managed-account alias rename submenu. + OpenManagedAccountRenamePanel, + /// Mark a managed account as active in the fork-owned registry. + SetManagedAccountActive(String), + /// Open an alias editor for a managed account. + OpenRenameManagedAccountAliasPrompt { + account_id: String, + current_alias: String, + }, + /// Persist a new alias for a managed account. + SaveManagedAccountAlias { + account_id: String, + alias: String, + }, /// Open the agent picker for switching active threads. OpenAgentPicker, /// Switch the active thread to the selected agent. @@ -100,6 +118,8 @@ pub(crate) enum AppEvent { /// Open the resume picker inside the running TUI session. OpenResumePicker, + /// Open the resume picker across all saved sessions, ignoring cwd filtering. + OpenResumePickerAll, /// Fork the current session into a new thread. ForkCurrentSession, diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 3729778b1..b8d13a3e1 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -308,9 +308,21 @@ use crate::streaming::controller::PlanStreamController; use crate::streaming::controller::StreamController; use chrono::Local; +use chrono::Utc; use codex_core::AuthManager; use codex_core::CodexAuth; use codex_core::ThreadManager; +use codex_ext::AccountManagementProfile; +use codex_ext::AccountPoolStore; +use codex_ext::AccountRateLimitSnapshot; +use codex_ext::AccountRateLimitWindow; +use codex_ext::DefaultAccountRouter; +use codex_ext::LimitSignalKind; +use codex_ext::RouteTurnRequest; +use codex_ext::RoutingTrigger; +use codex_ext::activate_managed_account; +use codex_ext::infer_limit_signal; +use codex_ext::persist_current_managed_account_snapshot; use codex_file_search::FileMatch; use codex_protocol::openai_models::InputModality; use codex_protocol::openai_models::ModelPreset; @@ -546,6 +558,47 @@ fn rate_limit_error_kind(info: &CodexErrorInfo) -> Option { } } +fn account_rate_limit_snapshot(snapshot: &RateLimitSnapshot) -> AccountRateLimitSnapshot { + AccountRateLimitSnapshot { + limit_name: snapshot.limit_name.clone(), + primary: snapshot + .primary + .as_ref() + .map(account_rate_limit_window_from_protocol), + secondary: snapshot + .secondary + .as_ref() + .map(account_rate_limit_window_from_protocol), + } +} + +fn account_rate_limit_snapshot_ref(snapshot: &RateLimitSnapshot) -> AccountRateLimitSnapshot { + account_rate_limit_snapshot(snapshot) +} + +fn account_rate_limit_window_from_protocol( + window: &codex_protocol::protocol::RateLimitWindow, +) -> AccountRateLimitWindow { + AccountRateLimitWindow { + used_percent: window.used_percent, + window_minutes: window.window_minutes, + resets_at: window.resets_at, + } +} + +fn mask_email(email: &str) -> String { + let Some((local, domain)) = email.split_once('@') else { + return email.to_string(); + }; + let visible = local.chars().take(3).collect::(); + let masked = if local.chars().count() > 3 { + format!("{visible}***") + } else { + format!("{visible}*") + }; + format!("{masked}@{domain}") +} + #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub(crate) enum ExternalEditorState { #[default] @@ -684,6 +737,9 @@ pub(crate) struct ChatWidget { initial_user_message: Option, token_info: Option, rate_limit_snapshots_by_limit_id: BTreeMap, + latest_codex_rate_limit_snapshot: Option, + last_submitted_user_turn: Option, + managed_account_retry_attempted: bool, plan_type: Option, rate_limit_warnings: RateLimitWarningState, rate_limit_switch_prompt: RateLimitSwitchPromptState, @@ -891,6 +947,12 @@ pub(crate) struct UserMessage { mention_bindings: Vec, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum UserTurnSubmissionKind { + Normal, + AutoRetry, +} + #[derive(Debug, Clone, PartialEq, Default)] struct ThreadComposerState { text: String, @@ -1749,6 +1811,7 @@ impl ChatWidget { fn on_task_complete(&mut self, last_agent_message: Option, from_replay: bool) { self.submit_pending_steers_after_interrupt = false; + self.managed_account_retry_attempted = false; if let Some(message) = last_agent_message.as_ref() && !message.trim().is_empty() { @@ -2022,6 +2085,10 @@ impl ChatWidget { self.plan_type = snapshot.plan_type.or(self.plan_type); let is_codex_limit = limit_id.eq_ignore_ascii_case("codex"); + if is_codex_limit { + self.latest_codex_rate_limit_snapshot = Some(snapshot.clone()); + self.persist_managed_account_rate_limit_snapshot(&snapshot); + } let warnings = if is_codex_limit { self.rate_limit_warnings.take_warnings( snapshot @@ -2078,6 +2145,7 @@ impl ChatWidget { } } else { self.rate_limit_snapshots_by_limit_id.clear(); + self.latest_codex_rate_limit_snapshot = None; } self.refresh_status_surfaces(); } @@ -3647,6 +3715,9 @@ impl ChatWidget { initial_user_message, token_info: None, rate_limit_snapshots_by_limit_id: BTreeMap::new(), + latest_codex_rate_limit_snapshot: None, + last_submitted_user_turn: None, + managed_account_retry_attempted: false, plan_type: None, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), @@ -3849,6 +3920,9 @@ impl ChatWidget { initial_user_message, token_info: None, rate_limit_snapshots_by_limit_id: BTreeMap::new(), + latest_codex_rate_limit_snapshot: None, + last_submitted_user_turn: None, + managed_account_retry_attempted: false, plan_type: None, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), @@ -4043,6 +4117,9 @@ impl ChatWidget { initial_user_message, token_info: None, rate_limit_snapshots_by_limit_id: BTreeMap::new(), + latest_codex_rate_limit_snapshot: None, + last_submitted_user_turn: None, + managed_account_retry_attempted: false, plan_type: None, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), @@ -4164,6 +4241,15 @@ impl ChatWidget { pub(crate) fn handle_key_event(&mut self, key_event: KeyEvent) { match key_event { + KeyEvent { + code: KeyCode::Char(c), + modifiers, + kind: KeyEventKind::Press, + .. + } if modifiers.contains(KeyModifiers::CONTROL) && c.eq_ignore_ascii_case(&'p') => { + self.app_event_tx.send(AppEvent::OpenControlPanel); + return; + } KeyEvent { code: KeyCode::Char(c), modifiers, @@ -4903,6 +4989,250 @@ impl ChatWidget { self.bottom_pane.show_view(Box::new(view)); } + fn current_managed_account_profile(&self) -> Option { + let auth = self.auth_manager.auth_cached().or_else(|| { + CodexAuth::from_auth_storage( + &self.config.codex_home, + self.config.cli_auth_credentials_store_mode, + ) + .ok() + .flatten() + })?; + if !auth.is_chatgpt_auth() { + return None; + } + + let id = auth.get_account_id()?; + Some(AccountManagementProfile { + id: id.clone(), + alias: Some(id), + masked_email: auth.get_account_email().map(|email| mask_email(&email)), + plan_label: auth + .account_plan_type() + .map(|plan_type| format!("{plan_type:?}").to_ascii_lowercase()), + priority: None, + }) + } + + fn persist_managed_account_rate_limit_snapshot(&self, snapshot: &RateLimitSnapshot) { + let Some(mut profile) = self.current_managed_account_profile() else { + return; + }; + if let Some(plan_type) = snapshot.plan_type { + profile.plan_label = Some(format!("{plan_type:?}").to_ascii_lowercase()); + } + let store = AccountPoolStore::new(self.config.codex_home.clone()); + let account_id = profile.id.clone(); + let rate_limit_snapshot = account_rate_limit_snapshot(snapshot); + if let Err(err) = store.update(|state| { + state.upsert_account(profile); + state.apply_rate_limit_snapshot(&account_id, &rate_limit_snapshot); + }) { + tracing::warn!("failed to persist managed account rate limit snapshot: {err}"); + } + } + + fn persist_managed_account_limit_signal(&self, kind: LimitSignalKind) { + let Some(profile) = self.current_managed_account_profile() else { + return; + }; + let store = AccountPoolStore::new(self.config.codex_home.clone()); + let account_id = profile.id.clone(); + let snapshot = self + .latest_codex_rate_limit_snapshot + .as_ref() + .map(account_rate_limit_snapshot_ref); + let signal = infer_limit_signal(kind, Utc::now().timestamp(), snapshot.as_ref()); + if let Err(err) = store.update(|state| { + state.upsert_account(profile); + state.apply_limit_signal(&account_id, &signal); + }) { + tracing::warn!("failed to persist managed account limit signal: {err}"); + } + } + + fn prepare_managed_account_for_user_turn(&mut self) -> bool { + let current_snapshot = match persist_current_managed_account_snapshot( + &self.config.codex_home, + self.config.cli_auth_credentials_store_mode, + ) { + Ok(snapshot) => snapshot, + Err(err) => { + self.add_error_message(format!( + "Failed to snapshot the current managed account before sending: {err}" + )); + return false; + } + }; + let Some(current_snapshot) = current_snapshot else { + return true; + }; + let store = AccountPoolStore::new(self.config.codex_home.clone()); + let mut state = match store.load() { + Ok(state) => state, + Err(err) => { + self.add_error_message(format!("Failed to read managed account pool: {err}")); + return false; + } + }; + state.upsert_account(current_snapshot.profile.clone()); + if state.accounts.is_empty() { + return true; + } + + let now_ts = Utc::now().timestamp(); + let current_account_id = Some(current_snapshot.profile.id); + let decision = DefaultAccountRouter::default().select_account( + &state, + &RouteTurnRequest { + now_ts, + trigger: RoutingTrigger::NormalTurn, + active_account_id: current_account_id.clone(), + preferred_account_id: state.active_account_id.clone(), + }, + ); + let Some(target_account_id) = decision.account_id else { + self.add_error_message( + "No healthy managed ChatGPT account is available for the next turn.".to_string(), + ); + return false; + }; + + if current_account_id.as_deref() == Some(target_account_id.as_str()) { + return true; + } + + let target_label = state + .accounts + .iter() + .find(|account| account.id == target_account_id) + .map(|account| account.display_name().to_string()) + .unwrap_or_else(|| target_account_id.clone()); + if let Err(err) = activate_managed_account( + &self.config.codex_home, + self.config.cli_auth_credentials_store_mode, + &target_account_id, + ) { + self.add_error_message(format!( + "Failed to activate managed account {target_label}: {err}" + )); + return false; + } + self.auth_manager.reload(); + if let Err(err) = store.update(|state| { + state.set_active_account(&target_account_id, now_ts); + }) { + tracing::warn!("failed to update managed account selection after activation: {err}"); + } + self.add_info_message( + format!("Switched to managed account {target_label} before sending."), + /*hint*/ None, + ); + true + } + + fn retry_user_turn_with_managed_account(&mut self) -> bool { + if self.managed_account_retry_attempted { + return false; + } + let Some(user_message) = self.last_submitted_user_turn.clone() else { + return false; + }; + + let current_snapshot = match persist_current_managed_account_snapshot( + &self.config.codex_home, + self.config.cli_auth_credentials_store_mode, + ) { + Ok(snapshot) => snapshot, + Err(err) => { + tracing::warn!("failed to snapshot current managed account before retry: {err}"); + return false; + } + }; + let Some(current_snapshot) = current_snapshot else { + return false; + }; + let store = AccountPoolStore::new(self.config.codex_home.clone()); + let mut state = match store.load() { + Ok(state) => state, + Err(err) => { + tracing::warn!("failed to load managed account pool for retry: {err}"); + return false; + } + }; + state.upsert_account(current_snapshot.profile.clone()); + + let now_ts = Utc::now().timestamp(); + let current_account_id = Some(current_snapshot.profile.id); + let decision = DefaultAccountRouter::default().select_account( + &state, + &RouteTurnRequest { + now_ts, + trigger: RoutingTrigger::RetryAfterHardError, + active_account_id: current_account_id.clone(), + preferred_account_id: state.active_account_id.clone(), + }, + ); + let Some(target_account_id) = decision.account_id else { + return false; + }; + if current_account_id.as_deref() == Some(target_account_id.as_str()) { + return false; + } + + let target_label = state + .accounts + .iter() + .find(|account| account.id == target_account_id) + .map(|account| account.display_name().to_string()) + .unwrap_or_else(|| target_account_id.clone()); + if let Err(err) = activate_managed_account( + &self.config.codex_home, + self.config.cli_auth_credentials_store_mode, + &target_account_id, + ) { + tracing::warn!("failed to activate managed account {target_label}: {err}"); + return false; + } + self.auth_manager.reload(); + if let Err(err) = store.update(|state| { + state.set_active_account(&target_account_id, now_ts); + }) { + tracing::warn!("failed to update managed account selection after retry: {err}"); + } + + self.managed_account_retry_attempted = true; + self.submit_pending_steers_after_interrupt = false; + self.finalize_turn(); + self.add_to_history(history_cell::new_info_event( + format!("Retrying the last turn with managed account {target_label}."), + /*hint*/ None, + )); + self.submit_user_message_with_kind(user_message, UserTurnSubmissionKind::AutoRetry); + true + } + + pub(crate) fn open_managed_account_alias_prompt( + &mut self, + account_id: String, + current_alias: String, + ) { + let tx = self.app_event_tx.clone(); + let view = CustomPromptView::new( + "Rename managed account".to_string(), + "Type a new alias and press Enter".to_string(), + Some(format!("Current alias: {current_alias}")), + Box::new(move |alias: String| { + tx.send(AppEvent::SaveManagedAccountAlias { + account_id: account_id.clone(), + alias, + }); + }), + ); + + self.bottom_pane.show_view(Box::new(view)); + } + pub(crate) fn handle_paste(&mut self, text: String) { self.bottom_pane.handle_paste(text); } @@ -4966,6 +5296,14 @@ impl ChatWidget { } fn submit_user_message(&mut self, user_message: UserMessage) { + self.submit_user_message_with_kind(user_message, UserTurnSubmissionKind::Normal); + } + + fn submit_user_message_with_kind( + &mut self, + user_message: UserMessage, + submission_kind: UserTurnSubmissionKind, + ) { if !self.is_session_configured() { tracing::warn!("cannot submit user message before session is configured; queueing"); self.queued_user_messages.push_front(user_message); @@ -5001,7 +5339,21 @@ impl ChatWidget { return; } - let render_in_history = !self.agent_turn_running; + let is_normal_submission = submission_kind == UserTurnSubmissionKind::Normal; + let render_in_history = is_normal_submission && !self.agent_turn_running; + if is_normal_submission { + self.last_submitted_user_turn = Some(UserMessage { + text: text.clone(), + local_images: local_images.clone(), + remote_image_urls: remote_image_urls.clone(), + text_elements: text_elements.clone(), + mention_bindings: mention_bindings.clone(), + }); + self.managed_account_retry_attempted = false; + if !self.prepare_managed_account_for_user_turn() { + return; + } + } let mut items: Vec = Vec::new(); // Special-case: "!cmd" executes a local shell command instead of sending to the model. @@ -5155,7 +5507,7 @@ impl ChatWidget { } else { None }; - let pending_steer = (!render_in_history).then(|| PendingSteer { + let pending_steer = (is_normal_submission && !render_in_history).then(|| PendingSteer { user_message: UserMessage { text: text.clone(), local_images: local_images.clone(), @@ -5190,7 +5542,7 @@ impl ChatWidget { } // Persist the text to cross-session message history. - if !text.is_empty() { + if is_normal_submission && !text.is_empty() { let encoded_mentions = mention_bindings .iter() .map(|binding| LinkedMention { @@ -5213,7 +5565,7 @@ impl ChatWidget { } // Show replayable user content in conversation history. - if render_in_history && !text.is_empty() { + if is_normal_submission && render_in_history && !text.is_empty() { let local_image_paths = local_images .into_iter() .map(|img| img.path) @@ -5231,7 +5583,7 @@ impl ChatWidget { local_image_paths, remote_image_urls, )); - } else if render_in_history && !remote_image_urls.is_empty() { + } else if is_normal_submission && render_in_history && !remote_image_urls.is_empty() { self.last_rendered_user_message_event = Some(Self::rendered_user_message_event_from_parts( String::new(), @@ -5402,7 +5754,14 @@ impl ChatWidget { self.on_server_overloaded_error(message) } RateLimitErrorKind::UsageLimit | RateLimitErrorKind::Generic => { - self.on_error(message) + self.persist_managed_account_limit_signal(match kind { + RateLimitErrorKind::UsageLimit => LimitSignalKind::UsageLimit, + RateLimitErrorKind::Generic => LimitSignalKind::RateLimit, + RateLimitErrorKind::ServerOverloaded => unreachable!(), + }); + if !self.retry_user_turn_with_managed_account() { + self.on_error(message); + } } } } else { diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index d837bb8ab..dd75a3c0a 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -36,6 +36,7 @@ use codex_core::config_loader::RequirementSource; use codex_core::models_manager::collaboration_mode_presets::CollaborationModesConfig; use codex_core::models_manager::manager::ModelsManager; use codex_core::skills::model::SkillMetadata; +use codex_ext::AccountPoolStore; use codex_features::FEATURES; use codex_features::Feature; use codex_otel::RuntimeMetricsSummary; @@ -1870,6 +1871,9 @@ async fn make_chatwidget_manual( initial_user_message: None, token_info: None, rate_limit_snapshots_by_limit_id: BTreeMap::new(), + latest_codex_rate_limit_snapshot: None, + last_submitted_user_turn: None, + managed_account_retry_attempted: false, plan_type: None, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), @@ -2470,6 +2474,80 @@ async fn rate_limit_switch_prompt_popup_snapshot() { assert_snapshot!("rate_limit_switch_prompt_popup", popup); } +#[tokio::test] +async fn rate_limit_snapshot_persists_managed_account_pool_state() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5")).await; + let codex_home = tempdir().expect("tempdir"); + chat.config.codex_home = codex_home.path().to_path_buf(); + chat.auth_manager = codex_core::test_support::auth_manager_from_auth( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + ); + let future_reset = chrono::Utc::now().timestamp() + 500; + + chat.on_rate_limit_snapshot(Some(RateLimitSnapshot { + limit_id: Some("codex".to_string()), + limit_name: Some("codex".to_string()), + primary: Some(RateLimitWindow { + used_percent: 87.0, + window_minutes: Some(300), + resets_at: Some(future_reset), + }), + secondary: None, + credits: None, + plan_type: Some(PlanType::Pro), + })); + + let state = AccountPoolStore::new(chat.config.codex_home.clone()) + .load() + .expect("load account pool"); + assert_eq!(state.active_account_id, Some("account_id".to_string())); + assert_eq!(state.accounts.len(), 1); + assert_eq!(state.accounts[0].plan_label, Some("pro".to_string())); + assert_eq!( + state.accounts[0].usage_summary(), + Some("5h 87/100".to_string()) + ); +} + +#[tokio::test] +async fn usage_limit_error_persists_managed_account_cooldown() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5")).await; + let codex_home = tempdir().expect("tempdir"); + chat.config.codex_home = codex_home.path().to_path_buf(); + chat.auth_manager = codex_core::test_support::auth_manager_from_auth( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + ); + let future_reset = chrono::Utc::now().timestamp() + 900; + let snapshot = RateLimitSnapshot { + limit_id: Some("codex".to_string()), + limit_name: Some("codex".to_string()), + primary: Some(RateLimitWindow { + used_percent: 100.0, + window_minutes: Some(300), + resets_at: Some(future_reset), + }), + secondary: None, + credits: None, + plan_type: Some(PlanType::Pro), + }; + chat.on_rate_limit_snapshot(Some(snapshot)); + + chat.dispatch_event_msg( + Some("evt-1".to_string()), + EventMsg::Error(ErrorEvent { + message: "You've hit your usage limit.".to_string(), + codex_error_info: Some(CodexErrorInfo::UsageLimitExceeded), + }), + None, + ); + + let state = AccountPoolStore::new(chat.config.codex_home.clone()) + .load() + .expect("load account pool"); + assert!(state.accounts[0].last_limit_error_at.is_some()); + assert_eq!(state.accounts[0].cooldown_until, Some(future_reset)); +} + #[tokio::test] async fn plan_implementation_popup_snapshot() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5")).await; @@ -6210,6 +6288,15 @@ async fn slash_fork_requests_current_fork() { assert_matches!(rx.try_recv(), Ok(AppEvent::ForkCurrentSession)); } +#[tokio::test] +async fn ctrl_p_opens_control_panel() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; + + chat.handle_key_event(KeyEvent::new(KeyCode::Char('p'), KeyModifiers::CONTROL)); + + assert_matches!(rx.try_recv(), Ok(AppEvent::OpenControlPanel)); +} + #[tokio::test] async fn slash_rollout_displays_current_path() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; diff --git a/codex-rs/tui_app_server/Cargo.toml b/codex-rs/tui_app_server/Cargo.toml index 886604205..68e75a5d7 100644 --- a/codex-rs/tui_app_server/Cargo.toml +++ b/codex-rs/tui_app_server/Cargo.toml @@ -41,6 +41,7 @@ codex-chatgpt = { workspace = true } codex-client = { workspace = true } codex-cloud-requirements = { workspace = true } codex-core = { workspace = true } +codex-ext = { workspace = true } codex-features = { workspace = true } codex-feedback = { workspace = true } codex-file-search = { workspace = true } diff --git a/codex-rs/tui_app_server/src/app.rs b/codex-rs/tui_app_server/src/app.rs index d4569ae7a..ce58a41e1 100644 --- a/codex-rs/tui_app_server/src/app.rs +++ b/codex-rs/tui_app_server/src/app.rs @@ -46,6 +46,7 @@ use crate::tui; use crate::tui::TuiEvent; use crate::update_action::UpdateAction; use crate::version::CODEX_CLI_VERSION; +use chrono::Utc; use codex_ansi_escape::ansi_escape_line; use codex_app_server_client::AppServerRequestHandle; use codex_app_server_protocol::ClientRequest; @@ -78,6 +79,10 @@ use codex_core::models_manager::model_presets::HIDE_GPT_5_1_CODEX_MAX_MIGRATION_ use codex_core::models_manager::model_presets::HIDE_GPT5_1_MIGRATION_PROMPT_CONFIG; #[cfg(target_os = "windows")] use codex_core::windows_sandbox::WindowsSandboxLevelExt; +use codex_ext::AccountPoolStore; +use codex_ext::HostCapabilities; +use codex_ext::activate_managed_account; +use codex_ext::persist_current_managed_account_snapshot; use codex_features::Feature; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; @@ -2500,6 +2505,297 @@ impl App { self.sync_active_agent_label(); } + fn open_control_panel(&mut self) { + let capabilities = HostCapabilities::codex_mvp(); + let items = vec![ + SelectionItem { + name: "Sessions".to_string(), + description: Some("Resume or switch saved chats.".to_string()), + actions: vec![Box::new(|tx| tx.send(AppEvent::OpenResumePickerAll))], + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Fork Current Session".to_string(), + description: Some("Fork the current thread into a new session.".to_string()), + actions: vec![Box::new(|tx| tx.send(AppEvent::ForkCurrentSession))], + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Accounts".to_string(), + description: Some("Inspect the managed multi-account pool.".to_string()), + actions: vec![Box::new(|tx| tx.send(AppEvent::OpenAccountsPanel))], + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Undo Last User Message".to_string(), + description: Some( + "Planned MVP follow-up: restore the last sent input and roll back one turn." + .to_string(), + ), + is_disabled: true, + disabled_reason: Some("Not wired yet.".to_string()), + ..Default::default() + }, + ]; + + self.chat_widget.show_selection_view(SelectionViewParams { + view_id: Some("fork-control-panel"), + title: Some("Control Panel".to_string()), + subtitle: Some(format!( + "Fork extension host: {} negotiated MVP capabilities.", + capabilities.capabilities.len() + )), + footer_hint: Some(standard_popup_hint_line()), + items, + ..Default::default() + }); + } + + fn open_accounts_panel(&mut self) { + let store = AccountPoolStore::new(self.config.codex_home.clone()); + let path = store.path(); + let state = self.sync_current_auth_into_account_pool(&store); + let mut items = Vec::new(); + let subtitle = match &state { + Ok(state) => Some(format!( + "{} account(s) configured. Active: {}.", + state.accounts.len(), + state.active_account_id.as_deref().unwrap_or("none") + )), + Err(err) => Some(format!("Failed to read account pool: {err}")), + }; + + match state { + Ok(state) => { + for account in &state.accounts { + let mut description_parts = Vec::new(); + if let Some(masked_email) = &account.masked_email { + description_parts.push(masked_email.clone()); + } + if let Some(plan_label) = &account.plan_label { + description_parts.push(plan_label.clone()); + } + if let Some(usage_summary) = account.usage_summary() { + description_parts.push(usage_summary); + } + if let Some(cooldown_until) = account.cooldown_until { + description_parts.push(format!("cooldown until {cooldown_until}")); + } + + let account_id = account.id.clone(); + let search_value = format!( + "{} {} {}", + account.display_name(), + account.id, + account.masked_email.clone().unwrap_or_default() + ); + items.push(SelectionItem { + name: account.display_name().to_string(), + description: (!description_parts.is_empty()) + .then(|| description_parts.join(" · ")), + is_current: state.active_account_id.as_deref() == Some(account.id.as_str()), + is_disabled: !account.enabled, + disabled_reason: (!account.enabled).then(|| "Disabled".to_string()), + actions: vec![Box::new(move |tx| { + tx.send(AppEvent::SetManagedAccountActive(account_id.clone())); + })], + dismiss_on_select: true, + search_value: Some(search_value), + ..Default::default() + }); + } + if !state.accounts.is_empty() { + items.push(SelectionItem { + name: "Rename".to_string(), + description: Some( + "Open alias rename actions for managed accounts.".to_string(), + ), + actions: vec![Box::new(|tx| { + tx.send(AppEvent::OpenManagedAccountRenamePanel) + })], + dismiss_on_select: true, + ..Default::default() + }); + } + } + Err(err) => { + items.push(SelectionItem { + name: "Account pool unavailable".to_string(), + description: Some(err.to_string()), + is_disabled: true, + ..Default::default() + }); + } + } + + if items.is_empty() { + items.push(SelectionItem { + name: "No managed accounts yet".to_string(), + description: Some(format!( + "Run repeated ChatGPT login flows and persist them into {}.", + path.display() + )), + is_disabled: true, + ..Default::default() + }); + } + + self.chat_widget.show_selection_view(SelectionViewParams { + view_id: Some("fork-accounts-panel"), + title: Some("Accounts".to_string()), + subtitle, + footer_hint: Some(standard_popup_hint_line()), + footer_note: Some(Line::from(path.display().to_string()).dim()), + items, + ..Default::default() + }); + } + + fn open_managed_account_rename_panel(&mut self) { + let store = AccountPoolStore::new(self.config.codex_home.clone()); + let path = store.path(); + let state = self.sync_current_auth_into_account_pool(&store); + let mut items = Vec::new(); + let subtitle = match &state { + Ok(state) => Some(format!( + "Rename aliases for {} managed account(s).", + state.accounts.len() + )), + Err(err) => Some(format!("Failed to read account pool: {err}")), + }; + + match state { + Ok(state) => { + for account in &state.accounts { + let account_id = account.id.clone(); + let current_alias = account.alias.clone(); + let mut description_parts = Vec::new(); + if let Some(alias) = + (!current_alias.is_empty()).then_some(current_alias.clone()) + { + description_parts.push(format!("alias: {alias}")); + } + description_parts.push(account.id.clone()); + if let Some(masked_email) = &account.masked_email { + description_parts.push(masked_email.clone()); + } + items.push(SelectionItem { + name: format!("Rename {}", account.display_name()), + description: Some(description_parts.join(" · ")), + actions: vec![Box::new(move |tx| { + tx.send(AppEvent::OpenRenameManagedAccountAliasPrompt { + account_id: account_id.clone(), + current_alias: current_alias.clone(), + }); + })], + dismiss_on_select: true, + ..Default::default() + }); + } + } + Err(err) => { + items.push(SelectionItem { + name: "Account pool unavailable".to_string(), + description: Some(err.to_string()), + is_disabled: true, + ..Default::default() + }); + } + } + + if items.is_empty() { + items.push(SelectionItem { + name: "No managed accounts yet".to_string(), + description: Some(format!( + "Run repeated ChatGPT login flows and persist them into {}.", + path.display() + )), + is_disabled: true, + ..Default::default() + }); + } + + self.chat_widget.show_selection_view(SelectionViewParams { + view_id: Some("fork-account-rename-panel"), + title: Some("Rename".to_string()), + subtitle, + footer_hint: Some(standard_popup_hint_line()), + footer_note: Some(Line::from(path.display().to_string()).dim()), + items, + ..Default::default() + }); + } + + fn sync_current_auth_into_account_pool( + &self, + store: &AccountPoolStore, + ) -> std::io::Result { + let snapshot = persist_current_managed_account_snapshot( + &self.config.codex_home, + self.config.cli_auth_credentials_store_mode, + )?; + store.update(|state| { + if let Some(snapshot) = &snapshot { + state.upsert_account(snapshot.profile.clone()); + } + }) + } + + fn set_managed_account_active(&mut self, account_id: String) { + if let Err(err) = persist_current_managed_account_snapshot( + &self.config.codex_home, + self.config.cli_auth_credentials_store_mode, + ) { + self.chat_widget + .add_to_history(history_cell::new_error_event(format!( + "Failed to snapshot current managed account auth: {err}" + ))); + return; + } + if let Err(err) = activate_managed_account( + &self.config.codex_home, + self.config.cli_auth_credentials_store_mode, + &account_id, + ) { + self.chat_widget + .add_to_history(history_cell::new_error_event(format!( + "Failed to activate managed account auth: {err}" + ))); + return; + } + + let store = AccountPoolStore::new(self.config.codex_home.clone()); + if let Err(err) = store.update(|state| { + state.set_active_account(&account_id, Utc::now().timestamp()); + }) { + self.chat_widget + .add_to_history(history_cell::new_error_event(format!( + "Failed to update active managed account: {err}" + ))); + return; + } + + self.open_accounts_panel(); + } + + fn save_managed_account_alias(&mut self, account_id: String, alias: String) { + let store = AccountPoolStore::new(self.config.codex_home.clone()); + if let Err(err) = store.update(|state| { + state.rename_account_alias(&account_id, alias.clone()); + }) { + self.chat_widget + .add_to_history(history_cell::new_error_event(format!( + "Failed to save managed account alias: {err}" + ))); + return; + } + + self.open_managed_account_rename_panel(); + } + /// Marks a cached picker thread closed and recomputes the contextual footer label. /// /// Closing a thread is not the same as removing it: users can still inspect finished agent @@ -3226,6 +3522,28 @@ impl App { event: AppEvent, ) -> Result { match event { + AppEvent::OpenControlPanel => { + self.open_control_panel(); + } + AppEvent::OpenAccountsPanel => { + self.open_accounts_panel(); + } + AppEvent::OpenManagedAccountRenamePanel => { + self.open_managed_account_rename_panel(); + } + AppEvent::SetManagedAccountActive(account_id) => { + self.set_managed_account_active(account_id); + } + AppEvent::OpenRenameManagedAccountAliasPrompt { + account_id, + current_alias, + } => { + self.chat_widget + .open_managed_account_alias_prompt(account_id, current_alias); + } + AppEvent::SaveManagedAccountAlias { account_id, alias } => { + self.save_managed_account_alias(account_id, alias); + } AppEvent::NewSession => { self.start_fresh_session_with_summary_hint(tui, app_server) .await; @@ -3354,6 +3672,123 @@ impl App { // Leaving alt-screen may blank the inline viewport; force a redraw either way. tui.frame_requester().schedule_frame(); } + AppEvent::OpenResumePickerAll => { + let picker_app_server = match crate::start_app_server_for_picker( + &self.config, + &match self.remote_app_server_url.clone() { + Some(websocket_url) => crate::AppServerTarget::Remote(websocket_url), + None => crate::AppServerTarget::Embedded, + }, + ) + .await + { + Ok(app_server) => app_server, + Err(err) => { + self.chat_widget.add_error_message(format!( + "Failed to start app-server-backed session picker: {err}" + )); + return Ok(AppRunControl::Continue); + } + }; + match crate::resume_picker::run_resume_picker_with_app_server( + tui, + &self.config, + /*show_all*/ true, + picker_app_server, + ) + .await? + { + SessionSelection::Resume(target_session) => { + let current_cwd = self.config.cwd.clone(); + let resume_cwd = if self.remote_app_server_url.is_some() { + current_cwd.clone() + } else { + match crate::resolve_cwd_for_resume_or_fork( + tui, + &self.config, + ¤t_cwd, + target_session.thread_id, + target_session.path.as_deref(), + CwdPromptAction::Resume, + /*allow_prompt*/ true, + ) + .await? + { + crate::ResolveCwdOutcome::Continue(Some(cwd)) => cwd, + crate::ResolveCwdOutcome::Continue(None) => current_cwd.clone(), + crate::ResolveCwdOutcome::Exit => { + return Ok(AppRunControl::Exit(ExitReason::UserRequested)); + } + } + }; + let mut resume_config = match self + .rebuild_config_for_resume_or_fallback(¤t_cwd, resume_cwd) + .await + { + Ok(cfg) => cfg, + Err(err) => { + self.chat_widget.add_error_message(format!( + "Failed to rebuild configuration for resume: {err}" + )); + return Ok(AppRunControl::Continue); + } + }; + self.apply_runtime_policy_overrides(&mut resume_config); + let summary = session_summary( + self.chat_widget.token_usage(), + self.chat_widget.thread_id(), + self.chat_widget.thread_name(), + ); + match app_server + .resume_thread(resume_config.clone(), target_session.thread_id) + .await + { + Ok(resumed) => { + self.shutdown_current_thread(app_server).await; + self.config = resume_config; + tui.set_notification_method(self.config.tui_notification_method); + self.file_search.update_search_dir(self.config.cwd.clone()); + match self + .replace_chat_widget_with_app_server_thread(tui, resumed) + .await + { + Ok(()) => { + if let Some(summary) = summary { + let mut lines: Vec> = + vec![summary.usage_line.clone().into()]; + if let Some(command) = summary.resume_command { + let spans = vec![ + "To continue this session, run ".into(), + command.cyan(), + ]; + lines.push(spans.into()); + } + self.chat_widget.add_plain_history_lines(lines); + } + } + Err(err) => { + self.chat_widget.add_error_message(format!( + "Failed to attach to resumed app-server thread: {err}" + )); + } + } + } + Err(err) => { + let path_display = target_session.display_label(); + self.chat_widget.add_error_message(format!( + "Failed to resume session from {path_display}: {err}" + )); + } + } + } + SessionSelection::Exit + | SessionSelection::StartFresh + | SessionSelection::Fork(_) => {} + } + + // Leaving alt-screen may blank the inline viewport; force a redraw either way. + tui.frame_requester().schedule_frame(); + } AppEvent::ForkCurrentSession => { self.session_telemetry.counter( "codex.thread.fork", diff --git a/codex-rs/tui_app_server/src/app_event.rs b/codex-rs/tui_app_server/src/app_event.rs index a763410dd..e7f7fad0c 100644 --- a/codex-rs/tui_app_server/src/app_event.rs +++ b/codex-rs/tui_app_server/src/app_event.rs @@ -74,6 +74,24 @@ pub(crate) struct ConnectorsSnapshot { #[allow(clippy::large_enum_variant)] #[derive(Debug)] pub(crate) enum AppEvent { + /// Open the fork-owned control panel. + OpenControlPanel, + /// Open the account pool panel inside the control panel flow. + OpenAccountsPanel, + /// Open the managed-account alias rename submenu. + OpenManagedAccountRenamePanel, + /// Mark a managed account as active in the fork-owned registry. + SetManagedAccountActive(String), + /// Open an alias editor for a managed account. + OpenRenameManagedAccountAliasPrompt { + account_id: String, + current_alias: String, + }, + /// Persist a new alias for a managed account. + SaveManagedAccountAlias { + account_id: String, + alias: String, + }, /// Open the agent picker for switching active threads. OpenAgentPicker, /// Switch the active thread to the selected agent. @@ -100,6 +118,8 @@ pub(crate) enum AppEvent { /// Open the resume picker inside the running TUI session. OpenResumePicker, + /// Open the resume picker across all saved sessions, ignoring cwd filtering. + OpenResumePickerAll, /// Fork the current session into a new thread. ForkCurrentSession, diff --git a/codex-rs/tui_app_server/src/chatwidget.rs b/codex-rs/tui_app_server/src/chatwidget.rs index 82cd0d99b..07a7cee67 100644 --- a/codex-rs/tui_app_server/src/chatwidget.rs +++ b/codex-rs/tui_app_server/src/chatwidget.rs @@ -341,6 +341,14 @@ use crate::streaming::controller::PlanStreamController; use crate::streaming::controller::StreamController; use chrono::Local; +use chrono::Utc; +use codex_core::CodexAuth; +use codex_ext::AccountManagementProfile; +use codex_ext::AccountPoolStore; +use codex_ext::AccountRateLimitSnapshot; +use codex_ext::AccountRateLimitWindow; +use codex_ext::LimitSignalKind; +use codex_ext::infer_limit_signal; use codex_file_search::FileMatch; use codex_protocol::openai_models::InputModality; use codex_protocol::openai_models::ModelPreset; @@ -587,6 +595,47 @@ fn app_server_rate_limit_error_kind(info: &AppServerCodexErrorInfo) -> Option AccountRateLimitSnapshot { + AccountRateLimitSnapshot { + limit_name: snapshot.limit_name.clone(), + primary: snapshot + .primary + .as_ref() + .map(account_rate_limit_window_from_protocol), + secondary: snapshot + .secondary + .as_ref() + .map(account_rate_limit_window_from_protocol), + } +} + +fn account_rate_limit_snapshot_ref(snapshot: &RateLimitSnapshot) -> AccountRateLimitSnapshot { + account_rate_limit_snapshot(snapshot) +} + +fn account_rate_limit_window_from_protocol( + window: &codex_protocol::protocol::RateLimitWindow, +) -> AccountRateLimitWindow { + AccountRateLimitWindow { + used_percent: window.used_percent, + window_minutes: window.window_minutes, + resets_at: window.resets_at, + } +} + +fn mask_email(email: &str) -> String { + let Some((local, domain)) = email.split_once('@') else { + return email.to_string(); + }; + let visible = local.chars().take(3).collect::(); + let masked = if local.chars().count() > 3 { + format!("{visible}***") + } else { + format!("{visible}*") + }; + format!("{masked}@{domain}") +} + #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub(crate) enum ExternalEditorState { #[default] @@ -726,6 +775,7 @@ pub(crate) struct ChatWidget { status_account_display: Option, token_info: Option, rate_limit_snapshots_by_limit_id: BTreeMap, + latest_codex_rate_limit_snapshot: Option, plan_type: Option, rate_limit_warnings: RateLimitWarningState, rate_limit_switch_prompt: RateLimitSwitchPromptState, @@ -2386,6 +2436,10 @@ impl ChatWidget { self.plan_type = snapshot.plan_type.or(self.plan_type); let is_codex_limit = limit_id.eq_ignore_ascii_case("codex"); + if is_codex_limit { + self.latest_codex_rate_limit_snapshot = Some(snapshot.clone()); + self.persist_managed_account_rate_limit_snapshot(&snapshot); + } let warnings = if is_codex_limit { self.rate_limit_warnings.take_warnings( snapshot @@ -2442,6 +2496,7 @@ impl ChatWidget { } } else { self.rate_limit_snapshots_by_limit_id.clear(); + self.latest_codex_rate_limit_snapshot = None; } self.refresh_status_line(); } @@ -2506,6 +2561,11 @@ impl ChatWidget { match info { RateLimitErrorKind::ServerOverloaded => self.on_server_overloaded_error(message), RateLimitErrorKind::UsageLimit | RateLimitErrorKind::Generic => { + self.persist_managed_account_limit_signal(match info { + RateLimitErrorKind::UsageLimit => LimitSignalKind::UsageLimit, + RateLimitErrorKind::Generic => LimitSignalKind::RateLimit, + RateLimitErrorKind::ServerOverloaded => unreachable!(), + }); self.on_error(message) } } @@ -4201,6 +4261,7 @@ impl ChatWidget { status_account_display, token_info: None, rate_limit_snapshots_by_limit_id: BTreeMap::new(), + latest_codex_rate_limit_snapshot: None, plan_type: initial_plan_type, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), @@ -4314,6 +4375,15 @@ impl ChatWidget { pub(crate) fn handle_key_event(&mut self, key_event: KeyEvent) { match key_event { + KeyEvent { + code: KeyCode::Char(c), + modifiers, + kind: KeyEventKind::Press, + .. + } if modifiers.contains(KeyModifiers::CONTROL) && c.eq_ignore_ascii_case(&'p') => { + self.app_event_tx.send(AppEvent::OpenControlPanel); + return; + } KeyEvent { code: KeyCode::Char(c), modifiers, @@ -5040,6 +5110,87 @@ impl ChatWidget { self.bottom_pane.show_view(Box::new(view)); } + fn current_managed_account_profile(&self) -> Option { + let auth = CodexAuth::from_auth_storage( + &self.config.codex_home, + self.config.cli_auth_credentials_store_mode, + ) + .ok() + .flatten()?; + if !auth.is_chatgpt_auth() { + return None; + } + + let id = auth.get_account_id()?; + Some(AccountManagementProfile { + id: id.clone(), + alias: Some(id), + masked_email: auth.get_account_email().map(|email| mask_email(&email)), + plan_label: auth + .account_plan_type() + .map(|plan_type| format!("{plan_type:?}").to_ascii_lowercase()), + priority: None, + }) + } + + fn persist_managed_account_rate_limit_snapshot(&self, snapshot: &RateLimitSnapshot) { + let Some(mut profile) = self.current_managed_account_profile() else { + return; + }; + if let Some(plan_type) = snapshot.plan_type { + profile.plan_label = Some(format!("{plan_type:?}").to_ascii_lowercase()); + } + let store = AccountPoolStore::new(self.config.codex_home.clone()); + let account_id = profile.id.clone(); + let rate_limit_snapshot = account_rate_limit_snapshot(snapshot); + if let Err(err) = store.update(|state| { + state.upsert_account(profile); + state.apply_rate_limit_snapshot(&account_id, &rate_limit_snapshot); + }) { + tracing::warn!("failed to persist managed account rate limit snapshot: {err}"); + } + } + + fn persist_managed_account_limit_signal(&self, kind: LimitSignalKind) { + let Some(profile) = self.current_managed_account_profile() else { + return; + }; + let store = AccountPoolStore::new(self.config.codex_home.clone()); + let account_id = profile.id.clone(); + let snapshot = self + .latest_codex_rate_limit_snapshot + .as_ref() + .map(account_rate_limit_snapshot_ref); + let signal = infer_limit_signal(kind, Utc::now().timestamp(), snapshot.as_ref()); + if let Err(err) = store.update(|state| { + state.upsert_account(profile); + state.apply_limit_signal(&account_id, &signal); + }) { + tracing::warn!("failed to persist managed account limit signal: {err}"); + } + } + + pub(crate) fn open_managed_account_alias_prompt( + &mut self, + account_id: String, + current_alias: String, + ) { + let tx = self.app_event_tx.clone(); + let view = CustomPromptView::new( + "Rename managed account".to_string(), + "Type a new alias and press Enter".to_string(), + Some(format!("Current alias: {current_alias}")), + Box::new(move |alias: String| { + tx.send(AppEvent::SaveManagedAccountAlias { + account_id: account_id.clone(), + alias, + }); + }), + ); + + self.bottom_pane.show_view(Box::new(view)); + } + pub(crate) fn handle_paste(&mut self, text: String) { self.bottom_pane.handle_paste(text); } @@ -6418,6 +6569,11 @@ impl ChatWidget { self.on_server_overloaded_error(message) } RateLimitErrorKind::UsageLimit | RateLimitErrorKind::Generic => { + self.persist_managed_account_limit_signal(match kind { + RateLimitErrorKind::UsageLimit => LimitSignalKind::UsageLimit, + RateLimitErrorKind::Generic => LimitSignalKind::RateLimit, + RateLimitErrorKind::ServerOverloaded => unreachable!(), + }); self.on_error(message) } } diff --git a/codex-rs/tui_app_server/src/chatwidget/tests.rs b/codex-rs/tui_app_server/src/chatwidget/tests.rs index 639b57da0..6b9475090 100644 --- a/codex-rs/tui_app_server/src/chatwidget/tests.rs +++ b/codex-rs/tui_app_server/src/chatwidget/tests.rs @@ -1892,6 +1892,7 @@ async fn make_chatwidget_manual( status_account_display: None, token_info: None, rate_limit_snapshots_by_limit_id: BTreeMap::new(), + latest_codex_rate_limit_snapshot: None, plan_type: None, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), @@ -6841,6 +6842,15 @@ async fn slash_fork_requests_current_fork() { assert_matches!(rx.try_recv(), Ok(AppEvent::ForkCurrentSession)); } +#[tokio::test] +async fn ctrl_p_opens_control_panel() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; + + chat.handle_key_event(KeyEvent::new(KeyCode::Char('p'), KeyModifiers::CONTROL)); + + assert_matches!(rx.try_recv(), Ok(AppEvent::OpenControlPanel)); +} + #[tokio::test] async fn slash_rollout_displays_current_path() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; diff --git a/docs/fork-extension-mvp.md b/docs/fork-extension-mvp.md new file mode 100644 index 000000000..6bda86e18 --- /dev/null +++ b/docs/fork-extension-mvp.md @@ -0,0 +1,70 @@ +# Fork Extension MVP + +## Proposal + +Maintain a forked `codex` binary while moving fork-specific behavior behind a +host-owned extension layer. The host stays small and explicit; policy and +product logic move into extension-facing data models and, later, WASM plugins. + +The first fork-owned target is multi-account ChatGPT routing with a fullscreen +control panel entry point. The fork should prefer additive crates and thin +integration shims over invasive edits to `codex-core`, `codex-tui`, or +`codex-tui-app-server`. + +## Design + +### Principles + +- Keep fork logic in a dedicated workspace crate: `codex-ext`. +- Treat host/plugin compatibility as capability negotiation, not lockstep API. +- Keep plugins behind host-controlled APIs so existing approval and sandbox + policies remain authoritative. +- Use the TUI as a consumer of extension state, not the source of truth. +- Route account switching above the provider layer, but below UI-specific code. + +### MVP Scope + +1. Add `codex-ext` with: + - host capability negotiation types + - persisted account-pool state under `CODEX_HOME/accounts/account-pool.json` + - persisted per-account auth snapshots under `CODEX_HOME/accounts//auth.json` + - a default threshold-based account router model + - login-time account snapshotting and alias-aware account registration +2. Add a `Ctrl-P` control panel entry point in both TUI implementations. +3. Show account-pool state in a read-only popup so the fork has a visible, + testable operator surface before auth switching is wired into requests. + +### Deferred + +- WASM runtime and ABI host +- plugin-owned TUI layouts +- undo-last-user-message implementation + +## Phased Todos + +### Phase 1 + +- Land `codex-ext` data model crate. +- Expose `Ctrl-P` control panel in both TUIs. +- Render account pool status from disk. + +### Phase 2 + +- Introduce fork-owned account registry alongside current single-auth storage. +- Add active-account selector and account alias management. +- Record inferred limit/cooldown signals from model errors. + +### Phase 3 + +- Route normal user turns through the default account router. +- Retry one normal user turn on explicit limit/rejection errors. +- Add `codex login --auth ` so browser/device-code login snapshots the + previous root auth and stores the new account under `accounts//`. +- Add observable switch reasons in the control panel. + +### Phase 4 + +- Add a WASM host with capability negotiation. +- Re-express the default account router as a built-in plugin module. +- Add hook/interceptor points for `AppStart`, `SessionStart`, + `BeforeTurnStart`, `BeforeToolCall`, and `AfterToolCall`.