diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index 59ae93bf1..5d7572f46 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -819,6 +819,8 @@ "dylint_linting_5.0.0": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"assert_cmd\",\"req\":\"^2.0\"},{\"name\":\"cargo_metadata\",\"req\":\"^0.23\"},{\"features\":[\"config\"],\"name\":\"dylint_internal\",\"req\":\"=5.0.0\"},{\"name\":\"paste\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"rustc_version\",\"req\":\"^0.4\"},{\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"name\":\"serde\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"tempfile\",\"req\":\"^3.23\"},{\"name\":\"thiserror\",\"req\":\"^2.0\"},{\"name\":\"toml\",\"req\":\"^0.9\"},{\"kind\":\"build\",\"name\":\"toml\",\"req\":\"^0.9\"}],\"features\":{\"constituent\":[]}}", "dylint_testing_5.0.0": "{\"dependencies\":[{\"name\":\"anyhow\",\"req\":\"^1.0\"},{\"name\":\"cargo_metadata\",\"req\":\"^0.23\"},{\"name\":\"compiletest_rs\",\"req\":\"^0.11\"},{\"name\":\"dylint\",\"req\":\"=5.0.0\"},{\"name\":\"dylint_internal\",\"req\":\"=5.0.0\"},{\"name\":\"env_logger\",\"req\":\"^0.11\"},{\"name\":\"once_cell\",\"req\":\"^1.21\"},{\"name\":\"regex\",\"req\":\"^1.11\"},{\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"name\":\"tempfile\",\"req\":\"^3.23\"}],\"features\":{\"default\":[],\"deny_warnings\":[]}}", "dyn-clone_1.0.20": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"features\":[\"diff\"],\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.66\"}],\"features\":{}}", + "ed25519-dalek_2.2.0": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"blake2\",\"req\":\"^0.10\"},{\"features\":[\"html_reports\"],\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.5\"},{\"default_features\":false,\"features\":[\"digest\"],\"name\":\"curve25519-dalek\",\"req\":\"^4\"},{\"default_features\":false,\"features\":[\"digest\",\"rand_core\"],\"kind\":\"dev\",\"name\":\"curve25519-dalek\",\"req\":\"^4\"},{\"default_features\":false,\"name\":\"ed25519\",\"req\":\">=2.2, <2.3\"},{\"kind\":\"dev\",\"name\":\"hex\",\"req\":\"^0.4\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"default_features\":false,\"name\":\"merlin\",\"optional\":true,\"req\":\"^3\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.8\"},{\"default_features\":false,\"name\":\"rand_core\",\"optional\":true,\"req\":\"^0.6.4\"},{\"default_features\":false,\"kind\":\"dev\",\"name\":\"rand_core\",\"req\":\"^0.6.4\"},{\"default_features\":false,\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"sha2\",\"req\":\"^0.10\"},{\"kind\":\"dev\",\"name\":\"sha3\",\"req\":\"^0.10\"},{\"default_features\":false,\"name\":\"signature\",\"optional\":true,\"req\":\">=2.0, <2.3\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2.3.0\"},{\"kind\":\"dev\",\"name\":\"toml\",\"req\":\"^0.7\"},{\"default_features\":false,\"features\":[\"static_secrets\"],\"kind\":\"dev\",\"name\":\"x25519-dalek\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"zeroize\",\"optional\":true,\"req\":\"^1.5\"}],\"features\":{\"alloc\":[\"curve25519-dalek/alloc\",\"ed25519/alloc\",\"serde?/alloc\",\"zeroize/alloc\"],\"asm\":[\"sha2/asm\"],\"batch\":[\"alloc\",\"merlin\",\"rand_core\"],\"default\":[\"fast\",\"std\",\"zeroize\"],\"digest\":[\"signature/digest\"],\"fast\":[\"curve25519-dalek/precomputed-tables\"],\"hazmat\":[],\"legacy_compatibility\":[\"curve25519-dalek/legacy_compatibility\"],\"pem\":[\"alloc\",\"ed25519/pem\",\"pkcs8\"],\"pkcs8\":[\"ed25519/pkcs8\"],\"rand_core\":[\"dep:rand_core\"],\"serde\":[\"dep:serde\",\"ed25519/serde\"],\"std\":[\"alloc\",\"ed25519/std\",\"serde?/std\",\"sha2/std\"],\"zeroize\":[\"dep:zeroize\",\"curve25519-dalek/zeroize\"]}}", + "ed25519_2.2.3": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1\"},{\"features\":[\"rand_core\"],\"kind\":\"dev\",\"name\":\"ed25519-dalek\",\"req\":\"^2\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"name\":\"pkcs8\",\"optional\":true,\"req\":\"^0.10\"},{\"features\":[\"std\"],\"kind\":\"dev\",\"name\":\"rand_core\",\"req\":\"^0.6\"},{\"default_features\":false,\"features\":[\"signature\"],\"kind\":\"dev\",\"name\":\"ring-compat\",\"req\":\"^0.8\"},{\"default_features\":false,\"name\":\"serde\",\"optional\":true,\"req\":\"^1\"},{\"name\":\"serde_bytes\",\"optional\":true,\"req\":\"^0.11\"},{\"default_features\":false,\"name\":\"signature\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"zeroize\",\"optional\":true,\"req\":\"^1\"}],\"features\":{\"alloc\":[\"pkcs8?/alloc\"],\"default\":[\"std\"],\"pem\":[\"alloc\",\"pkcs8/pem\"],\"serde_bytes\":[\"serde\",\"dep:serde_bytes\"],\"std\":[\"pkcs8?/std\",\"signature/std\"]}}", "either_1.15.0": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"alloc\",\"derive\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0.95\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0.0\"}],\"features\":{\"default\":[\"std\"],\"std\":[],\"use_std\":[\"std\"]}}", "ena_0.14.3": "{\"dependencies\":[{\"name\":\"dogged\",\"optional\":true,\"req\":\"^0.2.0\"},{\"name\":\"log\",\"req\":\"^0.4\"}],\"features\":{\"bench\":[],\"persistent\":[\"dogged\"]}}", "encode_unicode_1.0.0": "{\"dependencies\":[{\"default_features\":false,\"name\":\"ascii\",\"optional\":true,\"req\":\"^1.0.0\"},{\"kind\":\"dev\",\"name\":\"lazy_static\",\"req\":\"^1.0\",\"target\":\"cfg(unix)\"},{\"features\":[\"https-native\"],\"kind\":\"dev\",\"name\":\"minreq\",\"req\":\"^2.6\"}],\"features\":{\"default\":[\"std\"],\"std\":[]}}", diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 3a99c2aa6..5e54b772c 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1907,6 +1907,7 @@ dependencies = [ "codex-api", "codex-app-server-protocol", "codex-apply-patch", + "codex-arg0", "codex-async-utils", "codex-code-mode", "codex-config", @@ -1958,6 +1959,7 @@ dependencies = [ "ctor 0.6.3", "dirs", "dunce", + "ed25519-dalek", "env-flags", "eventsource-stream", "futures", @@ -3655,6 +3657,7 @@ dependencies = [ "cfg-if", "cpufeatures", "curve25519-dalek-derive", + "digest", "fiat-crypto", "rustc_version", "subtle", @@ -4131,7 +4134,7 @@ dependencies = [ "libc", "option-ext", "redox_users 0.5.2", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4235,6 +4238,30 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" +[[package]] +name = "ed25519" +version = "2.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" +dependencies = [ + "pkcs8", + "signature", +] + +[[package]] +name = "ed25519-dalek" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" +dependencies = [ + "curve25519-dalek", + "ed25519", + "serde", + "sha2", + "subtle", + "zeroize", +] + [[package]] name = "either" version = "1.15.0" @@ -6733,7 +6760,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -11857,7 +11884,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index e4985e458..b57a464d5 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -227,6 +227,7 @@ diffy = "0.4.2" dirs = "6" dotenvy = "0.15.7" dunce = "1.0.4" +ed25519-dalek = { version = "2.2.0", features = ["pkcs8"] } encoding_rs = "0.8.35" env-flags = "0.1.1" env_logger = "0.11.9" diff --git a/codex-rs/app-server/src/transport/remote_control/tests.rs b/codex-rs/app-server/src/transport/remote_control/tests.rs index 21808a3a1..6dd6978fa 100644 --- a/codex-rs/app-server/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server/src/transport/remote_control/tests.rs @@ -96,6 +96,7 @@ fn remote_control_auth_dot_json(account_id: Option<&str>) -> AuthDotJson { account_id: account_id.map(str::to_string), }), last_refresh: Some(chrono::Utc::now()), + agent_identity: None, } } diff --git a/codex-rs/app-server/src/transport/remote_control/websocket.rs b/codex-rs/app-server/src/transport/remote_control/websocket.rs index 42924c2d3..d168ec0e1 100644 --- a/codex-rs/app-server/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server/src/transport/remote_control/websocket.rs @@ -971,6 +971,7 @@ mod tests { account_id: Some("account_id".to_string()), }), last_refresh: Some(Utc::now()), + agent_identity: None, } } diff --git a/codex-rs/app-server/tests/common/auth_fixtures.rs b/codex-rs/app-server/tests/common/auth_fixtures.rs index 99334f077..86f0fb456 100644 --- a/codex-rs/app-server/tests/common/auth_fixtures.rs +++ b/codex-rs/app-server/tests/common/auth_fixtures.rs @@ -163,6 +163,7 @@ pub fn write_chatgpt_auth( openai_api_key: None, tokens: Some(tokens), last_refresh, + agent_identity: None, }; save_auth(codex_home, &auth, cli_auth_credentials_store_mode).context("write auth.json") diff --git a/codex-rs/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index 57a27961a..dbe61524f 100644 --- a/codex-rs/app-server/tests/suite/v2/app_list.rs +++ b/codex-rs/app-server/tests/suite/v2/app_list.rs @@ -117,6 +117,7 @@ async fn list_apps_returns_empty_with_api_key_auth() -> Result<()> { openai_api_key: Some("test-api-key".to_string()), tokens: None, last_refresh: None, + agent_identity: None, }, AuthCredentialsStoreMode::File, )?; diff --git a/codex-rs/core/BUILD.bazel b/codex-rs/core/BUILD.bazel index ed0199627..434dc1f6a 100644 --- a/codex-rs/core/BUILD.bazel +++ b/codex-rs/core/BUILD.bazel @@ -53,6 +53,7 @@ codex_rust_crate( "//codex-rs/linux-sandbox:codex-linux-sandbox", "//codex-rs/rmcp-client:test_stdio_server", "//codex-rs/rmcp-client:test_streamable_http_server", + "//codex-rs/responses-api-proxy:codex-responses-api-proxy", "//codex-rs/cli:codex", ], ) diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index c91f07f70..a37196607 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -41,6 +41,7 @@ codex-login = { workspace = true } codex-mcp = { workspace = true } codex-model-provider-info = { workspace = true } codex-models-manager = { workspace = true } +ed25519-dalek = { workspace = true } codex-shell-command = { workspace = true } codex-execpolicy = { workspace = true } codex-git-utils = { workspace = true } @@ -143,6 +144,7 @@ codex-shell-escalation = { workspace = true } [dev-dependencies] assert_cmd = { workspace = true } assert_matches = { workspace = true } +codex-arg0 = { workspace = true } codex-otel = { workspace = true } codex-test-binary-support = { workspace = true } codex-utils-cargo-bin = { workspace = true } diff --git a/codex-rs/core/src/agent_identity.rs b/codex-rs/core/src/agent_identity.rs new file mode 100644 index 000000000..13a07cd96 --- /dev/null +++ b/codex-rs/core/src/agent_identity.rs @@ -0,0 +1,766 @@ +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use chrono::SecondsFormat; +use chrono::Utc; +use codex_features::Feature; +use codex_login::AgentIdentityAuthRecord; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::default_client::create_client; +use codex_protocol::protocol::SessionSource; +use ed25519_dalek::SigningKey; +use ed25519_dalek::VerifyingKey; +use ed25519_dalek::pkcs8::DecodePrivateKey; +use ed25519_dalek::pkcs8::EncodePrivateKey; +use rand::TryRngCore; +use rand::rngs::OsRng; +use serde::Deserialize; +use serde::Serialize; +use tokio::sync::Mutex; +use tracing::debug; +use tracing::info; +use tracing::warn; + +use crate::config::Config; + +const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); +const AGENT_IDENTITY_BISCUIT_TIMEOUT: Duration = Duration::from_secs(15); + +#[derive(Clone)] +pub(crate) struct AgentIdentityManager { + auth_manager: Arc, + chatgpt_base_url: String, + feature_enabled: bool, + abom: AgentBillOfMaterials, + ensure_lock: Arc>, +} + +impl std::fmt::Debug for AgentIdentityManager { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AgentIdentityManager") + .field("chatgpt_base_url", &self.chatgpt_base_url) + .field("feature_enabled", &self.feature_enabled) + .field("abom", &self.abom) + .finish_non_exhaustive() + } +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub(crate) struct StoredAgentIdentity { + pub(crate) binding_id: String, + pub(crate) chatgpt_account_id: String, + pub(crate) chatgpt_user_id: Option, + pub(crate) agent_runtime_id: String, + pub(crate) private_key_pkcs8_base64: String, + pub(crate) public_key_ssh: String, + pub(crate) registered_at: String, + pub(crate) abom: AgentBillOfMaterials, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub(crate) struct AgentBillOfMaterials { + pub(crate) agent_version: String, + pub(crate) agent_harness_id: String, + pub(crate) running_location: String, +} + +#[derive(Debug, Serialize)] +struct RegisterAgentRequest { + abom: AgentBillOfMaterials, + agent_public_key: String, + capabilities: Vec, +} + +#[derive(Debug, Deserialize)] +struct RegisterAgentResponse { + agent_runtime_id: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct AgentIdentityBinding { + binding_id: String, + chatgpt_account_id: String, + chatgpt_user_id: Option, + access_token: String, +} + +struct GeneratedAgentKeyMaterial { + private_key_pkcs8_base64: String, + public_key_ssh: String, +} + +impl AgentIdentityManager { + pub(crate) fn new( + config: &Config, + auth_manager: Arc, + session_source: SessionSource, + ) -> Self { + Self { + auth_manager, + chatgpt_base_url: config.chatgpt_base_url.clone(), + feature_enabled: config.features.enabled(Feature::UseAgentIdentity), + abom: build_abom(session_source), + ensure_lock: Arc::new(Mutex::new(())), + } + } + + pub(crate) fn is_enabled(&self) -> bool { + self.feature_enabled + } + + pub(crate) async fn ensure_registered_identity(&self) -> Result> { + if !self.feature_enabled { + return Ok(None); + } + + let Some(auth) = self.auth_manager.auth().await else { + debug!("skipping agent identity registration because no auth is available"); + return Ok(None); + }; + + let Some(binding) = + AgentIdentityBinding::from_auth(&auth, self.auth_manager.forced_chatgpt_workspace_id()) + else { + debug!("skipping agent identity registration because ChatGPT auth is unavailable"); + return Ok(None); + }; + + let _guard = self.ensure_lock.lock().await; + + if let Some(stored_identity) = self.load_stored_identity(&auth, &binding)? { + info!( + agent_runtime_id = %stored_identity.agent_runtime_id, + binding_id = %binding.binding_id, + "reusing stored agent identity" + ); + return Ok(Some(stored_identity)); + } + + let stored_identity = self.register_agent_identity(&binding).await?; + self.store_identity(&auth, &stored_identity)?; + Ok(Some(stored_identity)) + } + + async fn register_agent_identity( + &self, + binding: &AgentIdentityBinding, + ) -> Result { + let key_material = generate_agent_key_material()?; + let request_body = RegisterAgentRequest { + abom: self.abom.clone(), + agent_public_key: key_material.public_key_ssh.clone(), + capabilities: Vec::new(), + }; + + let url = agent_registration_url(&self.chatgpt_base_url); + let human_biscuit = self.mint_human_biscuit(binding, "POST", &url).await?; + let client = create_client(); + let response = client + .post(&url) + .header("X-OpenAI-Authorization", human_biscuit) + .json(&request_body) + .timeout(AGENT_REGISTRATION_TIMEOUT) + .send() + .await + .with_context(|| { + format!("failed to send agent identity registration request to {url}") + })?; + + if response.status().is_success() { + let response_body = response + .json::() + .await + .with_context(|| format!("failed to parse agent identity response from {url}"))?; + let stored_identity = StoredAgentIdentity { + binding_id: binding.binding_id.clone(), + chatgpt_account_id: binding.chatgpt_account_id.clone(), + chatgpt_user_id: binding.chatgpt_user_id.clone(), + agent_runtime_id: response_body.agent_runtime_id, + private_key_pkcs8_base64: key_material.private_key_pkcs8_base64, + public_key_ssh: key_material.public_key_ssh, + registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + abom: self.abom.clone(), + }; + info!( + agent_runtime_id = %stored_identity.agent_runtime_id, + binding_id = %binding.binding_id, + "registered agent identity" + ); + return Ok(stored_identity); + } + + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!("agent identity registration failed with status {status} from {url}: {body}") + } + + async fn mint_human_biscuit( + &self, + binding: &AgentIdentityBinding, + target_method: &str, + target_url: &str, + ) -> Result { + let url = agent_identity_biscuit_url(&self.chatgpt_base_url); + let request_id = agent_identity_request_id()?; + let client = create_client(); + let response = client + .get(&url) + .bearer_auth(&binding.access_token) + .header("X-Request-Id", request_id.clone()) + .header("X-Original-Method", target_method) + .header("X-Original-Url", target_url) + .timeout(AGENT_IDENTITY_BISCUIT_TIMEOUT) + .send() + .await + .with_context(|| format!("failed to send agent identity biscuit request to {url}"))?; + + if response.status().is_success() { + let human_biscuit = response + .headers() + .get("x-openai-authorization") + .context("agent identity biscuit response did not include x-openai-authorization")? + .to_str() + .context("agent identity biscuit response header was not valid UTF-8")? + .to_string(); + info!( + request_id = %request_id, + "minted human biscuit for agent identity registration" + ); + return Ok(human_biscuit); + } + + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!( + "agent identity biscuit minting failed with status {status} from {url}: {body}" + ) + } + + fn load_stored_identity( + &self, + auth: &CodexAuth, + binding: &AgentIdentityBinding, + ) -> Result> { + let Some(record) = auth.get_agent_identity(&binding.chatgpt_account_id) else { + return Ok(None); + }; + + let stored_identity = + match StoredAgentIdentity::from_auth_record(binding, record, self.abom.clone()) { + Ok(stored_identity) => stored_identity, + Err(error) => { + warn!( + binding_id = %binding.binding_id, + error = %error, + "stored agent identity is invalid; deleting cached value" + ); + auth.remove_agent_identity()?; + return Ok(None); + } + }; + + if !stored_identity.matches_binding(binding) { + warn!( + binding_id = %binding.binding_id, + "stored agent identity binding no longer matches current auth; deleting cached value" + ); + auth.remove_agent_identity()?; + return Ok(None); + } + + if let Err(error) = stored_identity.validate_key_material() { + warn!( + agent_runtime_id = %stored_identity.agent_runtime_id, + binding_id = %binding.binding_id, + error = %error, + "stored agent identity key material is invalid; deleting cached value" + ); + auth.remove_agent_identity()?; + return Ok(None); + } + + Ok(Some(stored_identity)) + } + + fn store_identity( + &self, + auth: &CodexAuth, + stored_identity: &StoredAgentIdentity, + ) -> Result<()> { + auth.set_agent_identity(stored_identity.to_auth_record())?; + Ok(()) + } + + #[cfg(test)] + fn new_for_tests( + auth_manager: Arc, + feature_enabled: bool, + chatgpt_base_url: String, + session_source: SessionSource, + ) -> Self { + Self { + auth_manager, + chatgpt_base_url, + feature_enabled, + abom: build_abom(session_source), + ensure_lock: Arc::new(Mutex::new(())), + } + } +} + +impl StoredAgentIdentity { + fn from_auth_record( + binding: &AgentIdentityBinding, + record: AgentIdentityAuthRecord, + abom: AgentBillOfMaterials, + ) -> Result { + if record.workspace_id != binding.chatgpt_account_id { + anyhow::bail!( + "stored agent identity workspace {:?} does not match current workspace {:?}", + record.workspace_id, + binding.chatgpt_account_id + ); + } + let signing_key = signing_key_from_private_key_pkcs8_base64(&record.agent_private_key)?; + Ok(Self { + binding_id: binding.binding_id.clone(), + chatgpt_account_id: binding.chatgpt_account_id.clone(), + chatgpt_user_id: record.chatgpt_user_id, + agent_runtime_id: record.agent_runtime_id, + private_key_pkcs8_base64: record.agent_private_key, + public_key_ssh: encode_ssh_ed25519_public_key(&signing_key.verifying_key()), + registered_at: record.registered_at, + abom, + }) + } + + fn to_auth_record(&self) -> AgentIdentityAuthRecord { + AgentIdentityAuthRecord { + workspace_id: self.chatgpt_account_id.clone(), + chatgpt_user_id: self.chatgpt_user_id.clone(), + agent_runtime_id: self.agent_runtime_id.clone(), + agent_private_key: self.private_key_pkcs8_base64.clone(), + registered_at: self.registered_at.clone(), + } + } + + fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool { + self.binding_id == binding.binding_id + && self.chatgpt_account_id == binding.chatgpt_account_id + && match binding.chatgpt_user_id.as_deref() { + Some(chatgpt_user_id) => self.chatgpt_user_id.as_deref() == Some(chatgpt_user_id), + None => true, + } + } + + fn validate_key_material(&self) -> Result<()> { + let signing_key = self.signing_key()?; + let derived_public_key = encode_ssh_ed25519_public_key(&signing_key.verifying_key()); + anyhow::ensure!( + self.public_key_ssh == derived_public_key, + "stored public key does not match the private key" + ); + Ok(()) + } + + pub(crate) fn signing_key(&self) -> Result { + signing_key_from_private_key_pkcs8_base64(&self.private_key_pkcs8_base64) + } +} + +impl AgentIdentityBinding { + fn from_auth(auth: &CodexAuth, forced_workspace_id: Option) -> Option { + if !auth.is_chatgpt_auth() { + return None; + } + + let token_data = auth.get_token_data().ok()?; + let resolved_account_id = + forced_workspace_id + .filter(|value| !value.is_empty()) + .or(token_data + .account_id + .clone() + .filter(|value| !value.is_empty()))?; + + Some(Self { + binding_id: format!("chatgpt-account-{resolved_account_id}"), + chatgpt_account_id: resolved_account_id, + chatgpt_user_id: token_data + .id_token + .chatgpt_user_id + .filter(|value| !value.is_empty()), + access_token: token_data.access_token, + }) + } +} + +fn build_abom(session_source: SessionSource) -> AgentBillOfMaterials { + AgentBillOfMaterials { + agent_version: env!("CARGO_PKG_VERSION").to_string(), + agent_harness_id: match &session_source { + SessionSource::VSCode => "codex-app".to_string(), + SessionSource::Cli + | SessionSource::Exec + | SessionSource::Mcp + | SessionSource::Custom(_) + | SessionSource::SubAgent(_) + | SessionSource::Unknown => "codex-cli".to_string(), + }, + running_location: format!("{}-{}", session_source, std::env::consts::OS), + } +} + +fn generate_agent_key_material() -> Result { + let mut secret_key_bytes = [0u8; 32]; + OsRng + .try_fill_bytes(&mut secret_key_bytes) + .context("failed to generate agent identity private key bytes")?; + let signing_key = SigningKey::from_bytes(&secret_key_bytes); + let private_key_pkcs8 = signing_key + .to_pkcs8_der() + .context("failed to encode agent identity private key as PKCS#8")?; + + Ok(GeneratedAgentKeyMaterial { + private_key_pkcs8_base64: BASE64_STANDARD.encode(private_key_pkcs8.as_bytes()), + public_key_ssh: encode_ssh_ed25519_public_key(&signing_key.verifying_key()), + }) +} + +fn encode_ssh_ed25519_public_key(verifying_key: &VerifyingKey) -> String { + let mut blob = Vec::with_capacity(4 + 11 + 4 + 32); + append_ssh_string(&mut blob, b"ssh-ed25519"); + append_ssh_string(&mut blob, verifying_key.as_bytes()); + format!("ssh-ed25519 {}", BASE64_STANDARD.encode(blob)) +} + +fn append_ssh_string(buf: &mut Vec, value: &[u8]) { + buf.extend_from_slice(&(value.len() as u32).to_be_bytes()); + buf.extend_from_slice(value); +} + +fn agent_registration_url(chatgpt_base_url: &str) -> String { + let trimmed = chatgpt_base_url.trim_end_matches('/'); + format!("{trimmed}/v1/agent/register") +} + +fn signing_key_from_private_key_pkcs8_base64(private_key_pkcs8_base64: &str) -> Result { + let private_key = BASE64_STANDARD + .decode(private_key_pkcs8_base64) + .context("stored agent identity private key is not valid base64")?; + SigningKey::from_pkcs8_der(&private_key) + .context("stored agent identity private key is not valid PKCS#8") +} + +fn agent_identity_biscuit_url(chatgpt_base_url: &str) -> String { + let trimmed = chatgpt_base_url.trim_end_matches('/'); + format!("{trimmed}/authenticate_app_v2") +} + +fn agent_identity_request_id() -> Result { + let mut request_id_bytes = [0u8; 16]; + OsRng + .try_fill_bytes(&mut request_id_bytes) + .context("failed to generate agent identity request id")?; + Ok(format!( + "codex-agent-identity-{}", + URL_SAFE_NO_PAD.encode(request_id_bytes) + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + use base64::engine::general_purpose::URL_SAFE_NO_PAD; + use codex_app_server_protocol::AuthMode as ApiAuthMode; + use codex_login::AuthCredentialsStoreMode; + use codex_login::AuthDotJson; + use codex_login::save_auth; + use codex_login::token_data::IdTokenInfo; + use codex_login::token_data::TokenData; + use pretty_assertions::assert_eq; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::header; + use wiremock::matchers::method; + use wiremock::matchers::path; + + #[tokio::test] + async fn ensure_registered_identity_skips_when_feature_is_disabled() { + let auth_manager = + AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ false, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + + assert_eq!(manager.ensure_registered_identity().await.unwrap(), None); + } + + #[tokio::test] + async fn ensure_registered_identity_skips_for_api_key_auth() { + let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-key")); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + + assert_eq!(manager.ensure_registered_identity().await.unwrap(), None); + } + + #[tokio::test] + async fn ensure_registered_identity_registers_and_reuses_cached_identity() { + let server = MockServer::start().await; + let chatgpt_base_url = server.uri(); + mount_human_biscuit(&server, &chatgpt_base_url).await; + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "agent_runtime_id": "agent_123", + }))) + .expect(1) + .mount(&server) + .await; + + let auth_manager = + AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + + let first = manager + .ensure_registered_identity() + .await + .unwrap() + .expect("identity should be registered"); + let second = manager + .ensure_registered_identity() + .await + .unwrap() + .expect("identity should be reused"); + + assert_eq!(first.agent_runtime_id, "agent_123"); + assert_eq!(first, second); + assert_eq!(first.abom.agent_harness_id, "codex-cli"); + assert_eq!(first.chatgpt_account_id, "account-123"); + assert_eq!(first.chatgpt_user_id.as_deref(), Some("user-123")); + } + + #[tokio::test] + async fn ensure_registered_identity_deletes_invalid_cached_identity_and_reregisters() { + let server = MockServer::start().await; + let chatgpt_base_url = server.uri(); + mount_human_biscuit(&server, &chatgpt_base_url).await; + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "agent_runtime_id": "agent_456", + }))) + .expect(1) + .mount(&server) + .await; + + let auth = make_chatgpt_auth("account-123", Some("user-123")); + let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + + let binding = + AgentIdentityBinding::from_auth(&auth, /*forced_workspace_id*/ None).expect("binding"); + auth.set_agent_identity(AgentIdentityAuthRecord { + workspace_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), + agent_runtime_id: "agent_invalid".to_string(), + agent_private_key: "not-valid-base64".to_string(), + registered_at: "2026-01-01T00:00:00Z".to_string(), + }) + .expect("seed invalid identity"); + + let stored = manager + .ensure_registered_identity() + .await + .unwrap() + .expect("identity should be registered"); + + assert_eq!(stored.agent_runtime_id, "agent_456"); + let persisted = auth + .get_agent_identity(&binding.chatgpt_account_id) + .expect("stored identity"); + assert_eq!(persisted.agent_runtime_id, "agent_456"); + } + + #[tokio::test] + async fn ensure_registered_identity_deletes_different_user_identity_and_reregisters() { + let server = MockServer::start().await; + let chatgpt_base_url = server.uri(); + mount_human_biscuit(&server, &chatgpt_base_url).await; + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "agent_runtime_id": "agent_new", + }))) + .expect(1) + .mount(&server) + .await; + + let auth = make_chatgpt_auth("account-123", Some("user-new")); + let stale_key = generate_agent_key_material().expect("key material"); + auth.set_agent_identity(AgentIdentityAuthRecord { + workspace_id: "account-123".to_string(), + chatgpt_user_id: Some("user-old".to_string()), + agent_runtime_id: "agent_old".to_string(), + agent_private_key: stale_key.private_key_pkcs8_base64, + registered_at: "2026-01-01T00:00:00Z".to_string(), + }) + .expect("seed stale identity"); + + let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + + let stored = manager + .ensure_registered_identity() + .await + .unwrap() + .expect("identity should be registered"); + + assert_eq!(stored.agent_runtime_id, "agent_new"); + assert_eq!(stored.chatgpt_user_id.as_deref(), Some("user-new")); + let persisted = auth + .get_agent_identity("account-123") + .expect("stored identity"); + assert_eq!(persisted.agent_runtime_id, "agent_new"); + assert_eq!(persisted.chatgpt_user_id.as_deref(), Some("user-new")); + } + + #[tokio::test] + async fn ensure_registered_identity_uses_chatgpt_base_url() { + let server = MockServer::start().await; + let chatgpt_base_url = format!("{}/backend-api", server.uri()); + mount_human_biscuit(&server, &chatgpt_base_url).await; + Mock::given(method("POST")) + .and(path("/backend-api/v1/agent/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "agent_runtime_id": "agent_canonical", + }))) + .expect(1) + .mount(&server) + .await; + + let auth_manager = + AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + + let stored = manager + .ensure_registered_identity() + .await + .unwrap() + .expect("identity should be registered"); + assert_eq!(stored.agent_runtime_id, "agent_canonical"); + } + + async fn mount_human_biscuit(server: &MockServer, chatgpt_base_url: &str) { + let biscuit_url = agent_identity_biscuit_url(chatgpt_base_url); + let biscuit_path = reqwest::Url::parse(&biscuit_url) + .expect("biscuit URL parses") + .path() + .to_string(); + let target_url = agent_registration_url(chatgpt_base_url); + Mock::given(method("GET")) + .and(path(biscuit_path)) + .and(header("authorization", "Bearer access-token-account-123")) + .and(header("x-original-method", "POST")) + .and(header("x-original-url", target_url)) + .respond_with( + ResponseTemplate::new(200).insert_header("x-openai-authorization", "human-biscuit"), + ) + .expect(1) + .mount(server) + .await; + } + + #[test] + fn encode_ssh_ed25519_public_key_matches_expected_wire_shape() { + let key_material = generate_agent_key_material().expect("key material"); + let (_, encoded_blob) = key_material + .public_key_ssh + .split_once(' ') + .expect("public key contains scheme"); + let decoded = BASE64_STANDARD.decode(encoded_blob).expect("base64"); + + assert_eq!(&decoded[..4], 11u32.to_be_bytes().as_slice()); + assert_eq!(&decoded[4..15], b"ssh-ed25519"); + assert_eq!(&decoded[15..19], 32u32.to_be_bytes().as_slice()); + assert_eq!(decoded.len(), 51); + } + + fn make_chatgpt_auth(account_id: &str, user_id: Option<&str>) -> CodexAuth { + let tempdir = tempfile::tempdir().expect("tempdir"); + let auth_json = AuthDotJson { + auth_mode: Some(ApiAuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + id_token: IdTokenInfo { + email: None, + chatgpt_plan_type: None, + chatgpt_user_id: user_id.map(ToOwned::to_owned), + chatgpt_account_id: Some(account_id.to_string()), + raw_jwt: fake_id_token(account_id, user_id), + }, + access_token: format!("access-token-{account_id}"), + refresh_token: "refresh-token".to_string(), + account_id: Some(account_id.to_string()), + }), + last_refresh: Some(Utc::now()), + agent_identity: None, + }; + save_auth(tempdir.path(), &auth_json, AuthCredentialsStoreMode::File).expect("save auth"); + CodexAuth::from_auth_storage(tempdir.path(), AuthCredentialsStoreMode::File) + .expect("load auth") + .expect("auth") + } + + fn fake_id_token(account_id: &str, user_id: Option<&str>) -> String { + let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#); + let payload = serde_json::json!({ + "https://api.openai.com/auth": { + "chatgpt_user_id": user_id, + "chatgpt_account_id": account_id, + } + }); + let payload = URL_SAFE_NO_PAD.encode(payload.to_string()); + format!("{header}.{payload}.signature") + } +} diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index c063a5c97..a3d759481 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -13,6 +13,7 @@ use crate::agent::Mailbox; use crate::agent::MailboxReceiver; use crate::agent::agent_status_from_event; use crate::agent::status::is_final; +use crate::agent_identity::AgentIdentityManager; use crate::apps::render_apps_section; use crate::commit_attribution::commit_message_trailer_instruction; use crate::compact; @@ -1510,6 +1511,56 @@ impl Session { }); } + fn start_agent_identity_registration(self: &Arc) { + if !self.services.agent_identity_manager.is_enabled() { + return; + } + + let weak_sess = Arc::downgrade(self); + let mut auth_state_rx = self.services.auth_manager.subscribe_auth_state(); + tokio::spawn(async move { + loop { + let Some(sess) = weak_sess.upgrade() else { + return; + }; + match sess + .services + .agent_identity_manager + .ensure_registered_identity() + .await + { + Ok(Some(_)) => return, + Ok(None) => { + drop(sess); + if auth_state_rx.changed().await.is_err() { + return; + } + } + Err(error) => { + sess.fail_agent_identity_registration(error).await; + return; + } + } + } + }); + } + + async fn fail_agent_identity_registration(self: &Arc, error: anyhow::Error) { + warn!(error = %error, "agent identity registration failed"); + let message = format!( + "Agent identity registration failed. Codex cannot continue while `features.use_agent_identity` is enabled: {error}" + ); + self.send_event_raw(Event { + id: self.next_internal_sub_id(), + msg: EventMsg::Error(ErrorEvent { + message, + codex_error_info: Some(CodexErrorInfo::Other), + }), + }) + .await; + handlers::shutdown(self, self.next_internal_sub_id()).await; + } + #[allow(clippy::too_many_arguments)] fn make_turn_context( conversation_id: ThreadId, @@ -2055,6 +2106,11 @@ impl Session { hooks, rollout: Mutex::new(rollout_recorder), user_shell: Arc::new(default_shell), + agent_identity_manager: Arc::new(AgentIdentityManager::new( + config.as_ref(), + Arc::clone(&auth_manager), + session_configuration.session_source.clone(), + )), shell_snapshot_tx, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -2152,6 +2208,7 @@ impl Session { // Start the watcher after SessionConfigured so it cannot emit earlier events. sess.start_skills_watcher_listener(); + sess.start_agent_identity_registration(); // Construct sandbox_state before MCP startup so it can be sent to each // MCP server immediately after it becomes ready (avoiding blocking). let sandbox_state = SandboxState { diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 0ff4ea5a2..46d4cda57 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -110,6 +110,7 @@ use opentelemetry::trace::TraceId; use std::path::Path; use std::time::Duration; use tokio::time::sleep; +use tokio::time::timeout; use tracing_opentelemetry::OpenTelemetrySpanExt; use codex_protocol::mcp::CallToolResult as McpCallToolResult; @@ -2814,6 +2815,11 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { }), rollout: Mutex::new(None), user_shell: Arc::new(default_user_shell()), + agent_identity_manager: Arc::new(crate::agent_identity::AgentIdentityManager::new( + config.as_ref(), + Arc::clone(&auth_manager), + session_configuration.session_source.clone(), + )), shell_snapshot_tx: watch::channel(None).0, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -3666,6 +3672,11 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( }), rollout: Mutex::new(None), user_shell: Arc::new(default_user_shell()), + agent_identity_manager: Arc::new(crate::agent_identity::AgentIdentityManager::new( + config.as_ref(), + Arc::clone(&auth_manager), + session_configuration.session_source.clone(), + )), shell_snapshot_tx: watch::channel(None).0, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -3770,6 +3781,42 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( make_session_and_context_with_dynamic_tools_and_rx(Vec::new()).await } +#[tokio::test] +async fn fail_agent_identity_registration_emits_error_and_shutdown() { + let (session, _turn_context, rx_event) = make_session_and_context_with_rx().await; + + session + .fail_agent_identity_registration(anyhow::anyhow!("registration exploded")) + .await; + + let error_event = timeout(Duration::from_secs(1), rx_event.recv()) + .await + .expect("error event should arrive") + .expect("error event should be readable"); + match error_event.msg { + EventMsg::Error(ErrorEvent { + message, + codex_error_info, + }) => { + assert_eq!( + message, + "Agent identity registration failed. Codex cannot continue while `features.use_agent_identity` is enabled: registration exploded".to_string() + ); + assert_eq!(codex_error_info, Some(CodexErrorInfo::Other)); + } + other => panic!("expected error event, got {other:?}"), + } + + let shutdown_event = timeout(Duration::from_secs(1), rx_event.recv()) + .await + .expect("shutdown event should arrive") + .expect("shutdown event should be readable"); + match shutdown_event.msg { + EventMsg::ShutdownComplete => {} + other => panic!("expected shutdown event, got {other:?}"), + } +} + #[tokio::test] async fn refresh_mcp_servers_is_deferred_until_next_turn() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 93d62faca..37b5dcf5d 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -5,6 +5,7 @@ // the TUI or the tracing stack). #![deny(clippy::print_stdout, clippy::print_stderr)] +mod agent_identity; mod apply_patch; mod apps; mod arc_monitor; diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 2e89585f9..4b6c5b30f 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use crate::RolloutRecorder; use crate::SkillsManager; use crate::agent::AgentControl; +use crate::agent_identity::AgentIdentityManager; use crate::client::ModelClient; use crate::config::StartedNetworkProxy; use crate::exec_policy::ExecPolicyManager; @@ -41,6 +42,7 @@ pub(crate) struct SessionServices { pub(crate) hooks: Hooks, pub(crate) rollout: Mutex>, pub(crate) user_shell: Arc, + pub(crate) agent_identity_manager: Arc, pub(crate) shell_snapshot_tx: watch::Sender>>, pub(crate) show_raw_agent_reasoning: bool, pub(crate) exec_policy: Arc, diff --git a/codex-rs/core/tests/suite/responses_api_proxy_headers.rs b/codex-rs/core/tests/suite/responses_api_proxy_headers.rs index 7dc888ffc..8199956ee 100644 --- a/codex-rs/core/tests/suite/responses_api_proxy_headers.rs +++ b/codex-rs/core/tests/suite/responses_api_proxy_headers.rs @@ -5,6 +5,11 @@ use anyhow::Result; use anyhow::anyhow; use codex_features::Feature; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::SandboxPolicy; +use codex_protocol::user_input::UserInput; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call; @@ -13,10 +18,12 @@ use core_test_support::responses::mount_sse_once_match; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; +use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; +use std::io::ErrorKind; use std::io::Write; use std::path::Path; use std::process::Child; @@ -29,8 +36,9 @@ use tempfile::TempDir; const PARENT_PROMPT: &str = "spawn a subagent and report when it is started"; const CHILD_PROMPT: &str = "child: say done"; const SPAWN_CALL_ID: &str = "spawn-call-1"; -const PROXY_START_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5); +const PROXY_START_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 30); const PROXY_POLL_INTERVAL: Duration = Duration::from_millis(/*millis*/ 20); +const TURN_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 60); struct ResponsesApiProxy { child: Child, @@ -40,8 +48,17 @@ struct ResponsesApiProxy { impl ResponsesApiProxy { fn start(upstream_url: &str, dump_dir: &Path) -> Result { let server_info = dump_dir.join("server-info.json"); - let mut child = StdCommand::new(codex_utils_cargo_bin::cargo_bin("codex")?) - .args(["responses-api-proxy", "--server-info"]) + let (proxy_program, use_codex_multitool) = + match codex_utils_cargo_bin::cargo_bin("codex-responses-api-proxy") { + Ok(path) => (path, false), + Err(_) => (codex_utils_cargo_bin::cargo_bin("codex")?, true), + }; + let mut command = StdCommand::new(proxy_program); + if use_codex_multitool { + command.arg("responses-api-proxy"); + } + let mut child = command + .args(["--server-info"]) .arg(&server_info) .args(["--upstream-url", upstream_url, "--dump-dir"]) .arg(dump_dir) @@ -58,15 +75,27 @@ impl ResponsesApiProxy { let deadline = Instant::now() + PROXY_START_TIMEOUT; loop { - if let Ok(info) = std::fs::read_to_string(&server_info) { - let port = serde_json::from_str::(&info)? - .get("port") - .and_then(Value::as_u64) - .ok_or_else(|| anyhow!("proxy server info missing port"))?; - return Ok(Self { - child, - port: u16::try_from(port)?, - }); + match std::fs::read_to_string(&server_info) { + Ok(info) => { + if !info.trim().is_empty() { + match serde_json::from_str::(&info) { + Ok(info) => { + let port = info + .get("port") + .and_then(Value::as_u64) + .ok_or_else(|| anyhow!("proxy server info missing port"))?; + return Ok(Self { + child, + port: u16::try_from(port)?, + }); + } + Err(err) if err.is_eof() => {} + Err(err) => return Err(err.into()), + } + } + } + Err(err) if err.kind() == ErrorKind::NotFound => {} + Err(err) => return Err(err.into()), } if let Some(status) = child.try_wait()? { return Err(anyhow!( @@ -144,7 +173,7 @@ async fn responses_api_proxy_dumps_parent_and_subagent_identity_headers() -> Res .expect("test config should allow feature update"); }); let test = builder.build(&server).await?; - test.submit_turn(PARENT_PROMPT).await?; + submit_turn_with_timeout(&test, PARENT_PROMPT, dump_dir.path()).await?; let dumps = wait_for_proxy_request_dumps(dump_dir.path())?; let parent = dumps @@ -178,6 +207,85 @@ async fn responses_api_proxy_dumps_parent_and_subagent_identity_headers() -> Res Ok(()) } +async fn submit_turn_with_timeout(test: &TestCodex, prompt: &str, dump_dir: &Path) -> Result<()> { + let session_model = test.session_configured.model.clone(); + test.codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: prompt.into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: test.config.cwd.to_path_buf(), + approval_policy: AskForApproval::OnRequest, + approvals_reviewer: None, + sandbox_policy: SandboxPolicy::WorkspaceWrite { + writable_roots: Vec::new(), + read_only_access: Default::default(), + network_access: false, + exclude_tmpdir_env_var: false, + exclude_slash_tmp: false, + }, + model: session_model, + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + }) + .await?; + + let turn_started = wait_for_event_result(test, "turn started", dump_dir, |event| { + matches!(event, EventMsg::TurnStarted(_)) + }) + .await?; + let EventMsg::TurnStarted(turn_started) = turn_started else { + unreachable!("event predicate only matches turn started events"); + }; + wait_for_event_result(test, "turn complete", dump_dir, |event| match event { + EventMsg::TurnComplete(event) => event.turn_id == turn_started.turn_id, + _ => false, + }) + .await?; + + Ok(()) +} + +async fn wait_for_event_result( + test: &TestCodex, + stage: &str, + dump_dir: &Path, + mut predicate: F, +) -> Result +where + F: FnMut(&EventMsg) -> bool, +{ + let mut seen_events = Vec::new(); + tokio::time::timeout(TURN_TIMEOUT, async { + loop { + let event = test.codex.next_event().await?; + seen_events.push(event_summary(&event.msg)); + if predicate(&event.msg) { + return Ok::(event.msg); + } + } + }) + .await + .map_err(|_| { + anyhow!( + "timed out waiting for {stage}; saw events: {}; {}", + seen_events.join(" | "), + proxy_dump_summary(dump_dir) + ) + })? +} + +fn event_summary(event: &EventMsg) -> String { + let mut summary = format!("{event:?}"); + summary.truncate(240); + summary +} + fn request_body_contains(req: &wiremock::Request, text: &str) -> bool { std::str::from_utf8(&req.body).is_ok_and(|body| body.contains(text)) } @@ -212,12 +320,54 @@ fn read_proxy_request_dumps(dump_dir: &Path) -> Result> { .and_then(|name| name.to_str()) .is_some_and(|name| name.ends_with("-request.json")) { - dumps.push(serde_json::from_str(&std::fs::read_to_string(&path)?)?); + let contents = std::fs::read_to_string(&path)?; + if contents.trim().is_empty() { + continue; + } + + match serde_json::from_str(&contents) { + Ok(dump) => dumps.push(dump), + Err(err) if err.is_eof() => continue, + Err(err) => return Err(err.into()), + } } } Ok(dumps) } +fn proxy_dump_summary(dump_dir: &Path) -> String { + match read_proxy_request_dumps(dump_dir) { + Ok(dumps) => { + let bodies = dumps + .iter() + .filter_map(|dump| dump.get("body")) + .map(Value::to_string) + .collect::>() + .join("; "); + format!("proxy wrote {} request dumps: {bodies}", dumps.len()) + } + Err(err) => format!("failed to read proxy request dumps: {err}"), + } +} + +#[test] +fn read_proxy_request_dumps_ignores_in_progress_files() -> Result<()> { + let dump_dir = TempDir::new()?; + std::fs::write(dump_dir.path().join("empty-request.json"), "")?; + std::fs::write(dump_dir.path().join("partial-request.json"), "{\"body\"")?; + std::fs::write( + dump_dir.path().join("complete-request.json"), + serde_json::to_string(&json!({ "body": "ready" }))?, + )?; + + assert_eq!( + read_proxy_request_dumps(dump_dir.path())?, + vec![json!({ "body": "ready" })] + ); + + Ok(()) +} + fn dump_body_contains(dump: &Value, text: &str) -> bool { dump.get("body") .is_some_and(|body| body.to_string().contains(text)) diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 27064b683..8a16f5939 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -16,6 +16,8 @@ use serde_json::json; use std::sync::Arc; use tempfile::TempDir; use tempfile::tempdir; +use tokio::time::Duration; +use tokio::time::timeout; #[tokio::test] async fn refresh_without_id_token() { @@ -135,6 +137,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() { account_id: None, }), last_refresh: Some(last_refresh), + agent_identity: None, }, auth_dot_json ); @@ -172,6 +175,7 @@ fn logout_removes_auth_file() -> Result<(), std::io::Error> { openai_api_key: Some("sk-test-key".to_string()), tokens: None, last_refresh: None, + agent_identity: None, }; super::save_auth(dir.path(), &auth_dot_json, AuthCredentialsStoreMode::File)?; let auth_file = get_auth_file(dir.path()); @@ -181,6 +185,49 @@ fn logout_removes_auth_file() -> Result<(), std::io::Error> { Ok(()) } +#[test] +fn chatgpt_auth_persists_agent_identity_for_workspace() { + let codex_home = tempdir().unwrap(); + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("account-123".to_string()), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + let auth = super::load_auth( + codex_home.path(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + ) + .expect("load auth") + .expect("auth available"); + let record = AgentIdentityAuthRecord { + workspace_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), + agent_runtime_id: "agent_123".to_string(), + agent_private_key: "pkcs8-base64".to_string(), + registered_at: "2026-04-13T12:00:00Z".to_string(), + }; + + auth.set_agent_identity(record.clone()) + .expect("set agent identity"); + + assert_eq!(auth.get_agent_identity("account-123"), Some(record.clone())); + assert_eq!(auth.get_agent_identity("other-account"), None); + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let persisted = storage + .load() + .expect("load auth") + .expect("auth should exist"); + assert_eq!(persisted.agent_identity, Some(record)); + + assert!(auth.remove_agent_identity().expect("remove agent identity")); + assert_eq!(auth.get_agent_identity("account-123"), None); +} + #[test] fn unauthorized_recovery_reports_mode_and_step_names() { let dir = tempdir().unwrap(); @@ -474,6 +521,67 @@ exit 1 } } +#[tokio::test] +async fn auth_manager_notifies_when_auth_state_changes() { + let dir = tempdir().unwrap(); + let manager = AuthManager::shared( + dir.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + ); + let mut auth_state_rx = manager.subscribe_auth_state(); + + save_auth( + dir.path(), + &AuthDotJson { + auth_mode: Some(ApiAuthMode::ApiKey), + openai_api_key: Some("sk-test-key".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + }, + AuthCredentialsStoreMode::File, + ) + .expect("save auth"); + + assert!( + manager.reload(), + "reload should report a changed auth state" + ); + timeout(Duration::from_secs(1), auth_state_rx.changed()) + .await + .expect("auth change notification should arrive") + .expect("auth state watch should remain open"); + + save_auth( + dir.path(), + &AuthDotJson { + auth_mode: Some(ApiAuthMode::ApiKey), + openai_api_key: Some("sk-updated-key".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + }, + AuthCredentialsStoreMode::File, + ) + .expect("save updated auth"); + + assert!( + !manager.reload(), + "reload remains mode-stable even when the underlying credentials change" + ); + timeout(Duration::from_secs(1), auth_state_rx.changed()) + .await + .expect("auth reload notification should still arrive") + .expect("auth state watch should remain open"); + + manager.set_forced_chatgpt_workspace_id(Some("workspace-123".to_string())); + timeout(Duration::from_secs(1), auth_state_rx.changed()) + .await + .expect("workspace change notification should arrive") + .expect("auth state watch should remain open"); +} + struct AuthFileParams { openai_api_key: Option, chatgpt_plan_type: Option, diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 15b129a16..8768171d9 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -13,6 +13,7 @@ use std::sync::Arc; use std::sync::Mutex; use std::sync::RwLock; use tokio::sync::Mutex as AsyncMutex; +use tokio::sync::watch; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::AuthMode as ApiAuthMode; @@ -20,6 +21,7 @@ use codex_protocol::config_types::ForcedLoginMethod; use codex_protocol::config_types::ModelProviderAuthInfo; use super::external_bearer::BearerTokenRefresher; +pub use crate::auth::storage::AgentIdentityAuthRecord; pub use crate::auth::storage::AuthDotJson; use crate::auth::storage::AuthStorageBackend; use crate::auth::storage::create_auth_storage; @@ -60,6 +62,7 @@ pub struct ChatgptAuth { #[derive(Debug, Clone)] pub struct ChatgptAuthTokens { state: ChatgptAuthState, + storage: Arc, } #[derive(Debug, Clone)] @@ -204,14 +207,13 @@ impl CodexAuth { client, }; + let storage = create_auth_storage(codex_home.to_path_buf(), storage_mode); match auth_mode { - ApiAuthMode::Chatgpt => { - let storage = create_auth_storage(codex_home.to_path_buf(), storage_mode); - Ok(Self::Chatgpt(ChatgptAuth { state, storage })) - } - ApiAuthMode::ChatgptAuthTokens => { - Ok(Self::ChatgptAuthTokens(ChatgptAuthTokens { state })) - } + ApiAuthMode::Chatgpt => Ok(Self::Chatgpt(ChatgptAuth { state, storage })), + ApiAuthMode::ChatgptAuthTokens => Ok(Self::ChatgptAuthTokens(ChatgptAuthTokens { + state, + storage, + })), ApiAuthMode::ApiKey => unreachable!("api key mode is handled above"), } } @@ -352,6 +354,52 @@ impl CodexAuth { self.get_current_auth_json().and_then(|t| t.tokens) } + pub fn get_agent_identity(&self, workspace_id: &str) -> Option { + self.get_current_auth_json() + .and_then(|auth| auth.agent_identity) + .filter(|identity| identity.workspace_id == workspace_id) + } + + pub fn set_agent_identity(&self, record: AgentIdentityAuthRecord) -> std::io::Result<()> { + let (state, storage) = match self { + Self::Chatgpt(auth) => (&auth.state, &auth.storage), + Self::ChatgptAuthTokens(auth) => (&auth.state, &auth.storage), + Self::ApiKey(_) => return Ok(()), + }; + let mut guard = state + .auth_dot_json + .lock() + .map_err(|_| std::io::Error::other("failed to lock auth state"))?; + let mut auth = guard + .clone() + .ok_or_else(|| std::io::Error::other("auth data is not available"))?; + auth.agent_identity = Some(record); + storage.save(&auth)?; + *guard = Some(auth); + Ok(()) + } + + pub fn remove_agent_identity(&self) -> std::io::Result { + let (state, storage) = match self { + Self::Chatgpt(auth) => (&auth.state, &auth.storage), + Self::ChatgptAuthTokens(auth) => (&auth.state, &auth.storage), + Self::ApiKey(_) => return Ok(false), + }; + let mut guard = state + .auth_dot_json + .lock() + .map_err(|_| std::io::Error::other("failed to lock auth state"))?; + let Some(mut auth) = guard.clone() else { + return Ok(false); + }; + let removed = auth.agent_identity.take().is_some(); + if removed { + storage.save(&auth)?; + *guard = Some(auth); + } + Ok(removed) + } + /// Consider this private to integration tests. pub fn create_dummy_chatgpt_auth_for_testing() -> Self { let auth_dot_json = AuthDotJson { @@ -364,6 +412,7 @@ impl CodexAuth { account_id: Some("account_id".to_string()), }), last_refresh: Some(Utc::now()), + agent_identity: None, }; let client = create_client(); @@ -439,6 +488,7 @@ pub fn login_with_api_key( openai_api_key: Some(api_key.to_string()), tokens: None, last_refresh: None, + agent_identity: None, }; save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode) } @@ -812,6 +862,7 @@ impl AuthDotJson { openai_api_key: None, tokens: Some(tokens), last_refresh: Some(Utc::now()), + agent_identity: None, }) } @@ -1107,6 +1158,7 @@ pub struct AuthManager { forced_chatgpt_workspace_id: RwLock>, refresh_lock: AsyncMutex<()>, external_auth: RwLock>>, + auth_state_tx: watch::Sender<()>, } /// Configuration view required to construct a shared [`AuthManager`]. @@ -1155,6 +1207,7 @@ impl AuthManager { enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> Self { + let (auth_state_tx, _) = watch::channel(()); let managed_auth = load_auth( &codex_home, enable_codex_api_key_env, @@ -1173,11 +1226,13 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: AsyncMutex::new(()), external_auth: RwLock::new(None), + auth_state_tx, } } /// Create an AuthManager with a specific CodexAuth, for testing only. pub fn from_auth_for_testing(auth: CodexAuth) -> Arc { + let (auth_state_tx, _) = watch::channel(()); let cached = CachedAuth { auth: Some(auth), permanent_refresh_failure: None, @@ -1191,11 +1246,13 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: AsyncMutex::new(()), external_auth: RwLock::new(None), + auth_state_tx, }) } /// Create an AuthManager with a specific CodexAuth and codex home, for testing only. pub fn from_auth_for_testing_with_home(auth: CodexAuth, codex_home: PathBuf) -> Arc { + let (auth_state_tx, _) = watch::channel(()); let cached = CachedAuth { auth: Some(auth), permanent_refresh_failure: None, @@ -1208,10 +1265,12 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: AsyncMutex::new(()), external_auth: RwLock::new(None), + auth_state_tx, }) } pub fn external_bearer_only(config: ModelProviderAuthInfo) -> Arc { + let (auth_state_tx, _) = watch::channel(()); Arc::new(Self { codex_home: PathBuf::from("non-existent"), inner: RwLock::new(CachedAuth { @@ -1225,6 +1284,7 @@ impl AuthManager { external_auth: RwLock::new(Some( Arc::new(BearerTokenRefresher::new(config)) as Arc )), + auth_state_tx, }) } @@ -1364,6 +1424,7 @@ impl AuthManager { } tracing::info!("Reloaded auth, changed: {changed}"); guard.auth = new_auth; + self.auth_state_tx.send_replace(()); changed } else { false @@ -1373,18 +1434,23 @@ impl AuthManager { pub fn set_external_auth(&self, external_auth: Arc) { if let Ok(mut guard) = self.external_auth.write() { *guard = Some(external_auth); + self.auth_state_tx.send_replace(()); } } pub fn clear_external_auth(&self) { if let Ok(mut guard) = self.external_auth.write() { *guard = None; + self.auth_state_tx.send_replace(()); } } pub fn set_forced_chatgpt_workspace_id(&self, workspace_id: Option) { - if let Ok(mut guard) = self.forced_chatgpt_workspace_id.write() { + if let Ok(mut guard) = self.forced_chatgpt_workspace_id.write() + && *guard != workspace_id + { *guard = workspace_id; + self.auth_state_tx.send_replace(()); } } @@ -1395,6 +1461,10 @@ impl AuthManager { .and_then(|guard| guard.clone()) } + pub fn subscribe_auth_state(&self) -> watch::Receiver<()> { + self.auth_state_tx.subscribe() + } + pub fn has_external_auth(&self) -> bool { self.external_auth().is_some() } @@ -1638,8 +1708,14 @@ impl AuthManager { ), ))); } - let auth_dot_json = + let mut auth_dot_json = AuthDotJson::from_external_tokens(&refreshed).map_err(RefreshTokenError::Transient)?; + if let Some(previous_auth) = self + .auth_cached() + .and_then(|auth| auth.get_current_auth_json()) + { + auth_dot_json.agent_identity = previous_auth.agent_identity; + } save_auth( &self.codex_home, &auth_dot_json, diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index 97e801415..f0d7a3169 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -39,6 +39,19 @@ pub struct AuthDotJson { #[serde(default, skip_serializing_if = "Option::is_none")] pub last_refresh: Option>, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub agent_identity: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] +pub struct AgentIdentityAuthRecord { + pub workspace_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub chatgpt_user_id: Option, + pub agent_runtime_id: String, + pub agent_private_key: String, + pub registered_at: String, } pub(super) fn get_auth_file(codex_home: &Path) -> PathBuf { diff --git a/codex-rs/login/src/auth/storage_tests.rs b/codex-rs/login/src/auth/storage_tests.rs index 4bf72c11b..2e1cc8502 100644 --- a/codex-rs/login/src/auth/storage_tests.rs +++ b/codex-rs/login/src/auth/storage_tests.rs @@ -18,6 +18,7 @@ async fn file_storage_load_returns_auth_dot_json() -> anyhow::Result<()> { openai_api_key: Some("test-key".to_string()), tokens: None, last_refresh: Some(Utc::now()), + agent_identity: None, }; storage @@ -38,6 +39,7 @@ async fn file_storage_save_persists_auth_dot_json() -> anyhow::Result<()> { openai_api_key: Some("test-key".to_string()), tokens: None, last_refresh: Some(Utc::now()), + agent_identity: None, }; let file = get_auth_file(codex_home.path()); @@ -52,6 +54,30 @@ async fn file_storage_save_persists_auth_dot_json() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn file_storage_persists_agent_identity() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let auth_dot_json = AuthDotJson { + auth_mode: Some(AuthMode::Chatgpt), + openai_api_key: None, + tokens: None, + last_refresh: Some(Utc::now()), + agent_identity: Some(AgentIdentityAuthRecord { + workspace_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), + agent_runtime_id: "agent_123".to_string(), + agent_private_key: "pkcs8-base64".to_string(), + registered_at: "2026-04-13T12:00:00Z".to_string(), + }), + }; + + storage.save(&auth_dot_json)?; + + assert_eq!(storage.load()?, Some(auth_dot_json)); + Ok(()) +} + #[test] fn file_storage_delete_removes_auth_file() -> anyhow::Result<()> { let dir = tempdir()?; @@ -60,6 +86,7 @@ fn file_storage_delete_removes_auth_file() -> anyhow::Result<()> { openai_api_key: Some("sk-test-key".to_string()), tokens: None, last_refresh: None, + agent_identity: None, }; let storage = create_auth_storage(dir.path().to_path_buf(), AuthCredentialsStoreMode::File); storage.save(&auth_dot_json)?; @@ -83,6 +110,7 @@ fn ephemeral_storage_save_load_delete_is_in_memory_only() -> anyhow::Result<()> openai_api_key: Some("sk-ephemeral".to_string()), tokens: None, last_refresh: Some(Utc::now()), + agent_identity: None, }; storage.save(&auth_dot_json)?; @@ -181,6 +209,7 @@ fn auth_with_prefix(prefix: &str) -> AuthDotJson { account_id: Some(format!("{prefix}-account-id")), }), last_refresh: None, + agent_identity: None, } } @@ -197,6 +226,7 @@ fn keyring_auth_storage_load_returns_deserialized_auth() -> anyhow::Result<()> { openai_api_key: Some("sk-test".to_string()), tokens: None, last_refresh: None, + agent_identity: None, }; seed_keyring_with_auth( &mock_keyring, @@ -239,6 +269,7 @@ fn keyring_auth_storage_save_persists_and_removes_fallback_file() -> anyhow::Res account_id: Some("account".to_string()), }), last_refresh: Some(Utc::now()), + agent_identity: None, }; storage.save(&auth)?; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 1b52223bc..4880e4310 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -20,6 +20,7 @@ pub use server::ShutdownHandle; pub use server::run_login_server; pub use api_bridge::auth_provider_from_auth; +pub use auth::AgentIdentityAuthRecord; pub use auth::AuthConfig; pub use auth::AuthDotJson; pub use auth::AuthManager; diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index 169a8a309..0c7b81018 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -781,6 +781,7 @@ pub(crate) async fn persist_tokens_async( openai_api_key: api_key, tokens: Some(tokens), last_refresh: Some(Utc::now()), + agent_identity: None, }; save_auth(&codex_home, &auth, auth_credentials_store_mode) }) diff --git a/codex-rs/login/tests/suite/auth_refresh.rs b/codex-rs/login/tests/suite/auth_refresh.rs index bf9e03bc2..d754c9589 100644 --- a/codex-rs/login/tests/suite/auth_refresh.rs +++ b/codex-rs/login/tests/suite/auth_refresh.rs @@ -54,6 +54,7 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -117,6 +118,7 @@ async fn refresh_token_refreshes_when_auth_is_unchanged() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -171,6 +173,7 @@ async fn refresh_token_skips_refresh_when_auth_changed() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -180,6 +183,7 @@ async fn refresh_token_skips_refresh_when_auth_changed() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -234,6 +238,7 @@ async fn refresh_token_errors_on_account_mismatch() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -244,6 +249,7 @@ async fn refresh_token_errors_on_account_mismatch() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -302,6 +308,7 @@ async fn returns_fresh_tokens_as_is() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(stale_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -349,6 +356,7 @@ async fn refreshes_token_when_access_token_is_expired() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(fresh_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -398,6 +406,7 @@ async fn auth_reloads_disk_auth_when_cached_auth_is_stale() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens), last_refresh: Some(stale_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -408,6 +417,7 @@ async fn auth_reloads_disk_auth_when_cached_auth_is_stale() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(fresh_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -459,6 +469,7 @@ async fn auth_reloads_disk_auth_without_calling_expired_refresh_token() -> Resul openai_api_key: None, tokens: Some(initial_tokens), last_refresh: Some(stale_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -469,6 +480,7 @@ async fn auth_reloads_disk_auth_without_calling_expired_refresh_token() -> Resul openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(fresh_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -518,6 +530,7 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -570,6 +583,7 @@ async fn refresh_token_does_not_retry_after_permanent_failure() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -636,6 +650,7 @@ async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result< openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -657,6 +672,7 @@ async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result< openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(fresh_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -715,6 +731,7 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -767,6 +784,7 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -776,6 +794,7 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -859,6 +878,7 @@ async fn unauthorized_recovery_errors_on_account_mismatch() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -869,6 +889,7 @@ async fn unauthorized_recovery_errors_on_account_mismatch() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -926,6 +947,7 @@ async fn unauthorized_recovery_requires_chatgpt_auth() -> Result<()> { openai_api_key: Some("sk-test".to_string()), tokens: None, last_refresh: None, + agent_identity: None, }; ctx.write_auth(&auth)?; diff --git a/codex-rs/tui/src/local_chatgpt_auth.rs b/codex-rs/tui/src/local_chatgpt_auth.rs index e888c0387..1f84b289a 100644 --- a/codex-rs/tui/src/local_chatgpt_auth.rs +++ b/codex-rs/tui/src/local_chatgpt_auth.rs @@ -108,6 +108,7 @@ mod tests { account_id: Some("workspace-1".to_string()), }), last_refresh: Some(Utc::now()), + agent_identity: None, }; save_auth(codex_home, &auth, AuthCredentialsStoreMode::File) .expect("chatgpt auth should save"); @@ -154,6 +155,7 @@ mod tests { openai_api_key: Some("sk-test".to_string()), tokens: None, last_refresh: None, + agent_identity: None, }, AuthCredentialsStoreMode::File, )