From fe4f6f94e0c4a65a3590d47dc27f027bbab668c5 Mon Sep 17 00:00:00 2001 From: adrian Date: Thu, 9 Apr 2026 12:14:32 -0700 Subject: [PATCH] Register agent identities behind use_agent_identity --- codex-rs/Cargo.lock | 33 +- codex-rs/Cargo.toml | 1 + codex-rs/core/Cargo.toml | 2 + codex-rs/core/src/agent_identity.rs | 685 ++++++++++++++++++++++++++++ codex-rs/core/src/codex.rs | 24 + codex-rs/core/src/codex_tests.rs | 10 + codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/state/service.rs | 2 + 8 files changed, 756 insertions(+), 3 deletions(-) create mode 100644 codex-rs/core/src/agent_identity.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 6c101b940d..0dd569e816 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1915,6 +1915,7 @@ dependencies = [ "codex-git-utils", "codex-hooks", "codex-instructions", + "codex-keyring-store", "codex-login", "codex-mcp", "codex-model-provider-info", @@ -1953,6 +1954,7 @@ dependencies = [ "ctor 0.6.3", "dirs", "dunce", + "ed25519-dalek", "env-flags", "eventsource-stream", "futures", @@ -3617,6 +3619,7 @@ dependencies = [ "cfg-if", "cpufeatures", "curve25519-dalek-derive", + "digest", "fiat-crypto", "rustc_version", "subtle", @@ -4087,7 +4090,7 @@ dependencies = [ "libc", "option-ext", "redox_users 0.5.2", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4191,6 +4194,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" @@ -6689,7 +6716,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]] @@ -11813,7 +11840,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 32ae50bfb7..4b5f253337 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -223,6 +223,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/core/Cargo.toml b/codex-rs/core/Cargo.toml index 55ce13afdc..539c474aae 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 } @@ -144,6 +145,7 @@ codex-shell-escalation = { workspace = true } assert_cmd = { workspace = true } assert_matches = { workspace = true } codex-arg0 = { workspace = true } +codex-keyring-store = { workspace = true } codex-otel = { workspace = true } codex-utils-cargo-bin = { workspace = true } core_test_support = { 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 0000000000..5ac40f2fe5 --- /dev/null +++ b/codex-rs/core/src/agent_identity.rs @@ -0,0 +1,685 @@ +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 chrono::SecondsFormat; +use chrono::Utc; +use codex_features::Feature; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_protocol::protocol::SessionSource; +use codex_secrets::SecretName; +use codex_secrets::SecretScope; +use codex_secrets::SecretsBackendKind; +use codex_secrets::SecretsManager; +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 reqwest::StatusCode; +use serde::Deserialize; +use serde::Serialize; +use tokio::sync::Mutex; +use tracing::debug; +use tracing::info; +use tracing::warn; + +use crate::config::Config; +use crate::default_client::create_client; + +const AGENT_IDENTITY_SECRET_NAME: &str = "AGENT_IDENTITY"; +const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); + +#[derive(Clone)] +pub(crate) struct AgentIdentityManager { + auth_manager: Arc, + secrets_manager: SecretsManager, + chatgpt_base_url: String, + feature_enabled: bool, + abom: AgentBillOfMaterials, + ensure_lock: Arc>, +} + +#[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, + secrets_manager: SecretsManager::new( + config.codex_home.clone(), + SecretsBackendKind::Local, + ), + 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) 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(&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(&binding, &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 client = create_client(); + let urls = agent_registration_urls(&self.chatgpt_base_url); + + for (index, url) in urls.iter().enumerate() { + let response = client + .post(url) + .bearer_auth(&binding.access_token) + .header("chatgpt-account-id", &binding.chatgpt_account_id) + .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(); + let is_last_candidate = index + 1 == urls.len(); + if !is_last_candidate + && matches!( + status, + StatusCode::NOT_FOUND | StatusCode::METHOD_NOT_ALLOWED + ) + { + debug!( + url = %url, + status = %status, + "agent identity registration endpoint unavailable at candidate URL; trying fallback" + ); + continue; + } + + anyhow::bail!( + "agent identity registration failed with status {status} from {url}: {body}" + ); + } + + anyhow::bail!("no candidate URLs were available for agent identity registration") + } + + fn load_stored_identity( + &self, + binding: &AgentIdentityBinding, + ) -> Result> { + let scope = secret_scope(binding)?; + let name = agent_identity_secret_name()?; + let Some(value) = self.secrets_manager.get(&scope, &name)? else { + return Ok(None); + }; + + let stored_identity = match serde_json::from_str::(&value) { + Ok(stored_identity) => stored_identity, + Err(error) => { + warn!( + binding_id = %binding.binding_id, + error = %error, + "stored agent identity could not be deserialized; deleting cached value" + ); + self.delete_identity(&scope, &name)?; + 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" + ); + self.delete_identity(&scope, &name)?; + 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" + ); + self.delete_identity(&scope, &name)?; + return Ok(None); + } + + Ok(Some(stored_identity)) + } + + fn store_identity( + &self, + binding: &AgentIdentityBinding, + stored_identity: &StoredAgentIdentity, + ) -> Result<()> { + let scope = secret_scope(binding)?; + let name = agent_identity_secret_name()?; + let value = serde_json::to_string(stored_identity) + .context("failed to serialize stored agent identity")?; + self.secrets_manager.set(&scope, &name, &value) + } + + fn delete_identity(&self, scope: &SecretScope, name: &SecretName) -> Result<()> { + self.secrets_manager.delete(scope, name)?; + Ok(()) + } + + #[cfg(test)] + fn new_for_tests( + auth_manager: Arc, + feature_enabled: bool, + chatgpt_base_url: String, + session_source: SessionSource, + secrets_manager: SecretsManager, + ) -> Self { + Self { + auth_manager, + secrets_manager, + chatgpt_base_url, + feature_enabled, + abom: build_abom(session_source), + ensure_lock: Arc::new(Mutex::new(())), + } + } +} + +impl StoredAgentIdentity { + 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 { + let private_key = BASE64_STANDARD + .decode(&self.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") + } +} + +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_identity_secret_name() -> Result { + SecretName::new(AGENT_IDENTITY_SECRET_NAME) + .context("agent identity secret name constant must be valid") +} + +fn secret_scope(binding: &AgentIdentityBinding) -> Result { + SecretScope::environment(binding.binding_id.clone()) + .context("agent identity binding must be a valid secrets scope") +} + +fn agent_registration_urls(chatgpt_base_url: &str) -> Vec { + let trimmed = chatgpt_base_url.trim_end_matches('/'); + if let Some(root) = trimmed.strip_suffix("/backend-api") { + return vec![ + format!("{root}/v1/agent/register"), + format!("{trimmed}/v1/agent/register"), + ]; + } + vec![format!("{trimmed}/v1/agent/register")] +} + +#[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_keyring_store::tests::MockKeyringStore; + 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 tempdir = tempfile::tempdir().expect("tempdir"); + let keyring_store = Arc::new(MockKeyringStore::default()); + let secrets_manager = SecretsManager::new_with_keyring_store( + tempdir.path().to_path_buf(), + SecretsBackendKind::Local, + keyring_store, + ); + 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, + secrets_manager, + ); + + assert_eq!(manager.ensure_registered_identity().await.unwrap(), None); + } + + #[tokio::test] + async fn ensure_registered_identity_skips_for_api_key_auth() { + let tempdir = tempfile::tempdir().expect("tempdir"); + let keyring_store = Arc::new(MockKeyringStore::default()); + let secrets_manager = SecretsManager::new_with_keyring_store( + tempdir.path().to_path_buf(), + SecretsBackendKind::Local, + keyring_store, + ); + 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, + secrets_manager, + ); + + 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; + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .and(header("authorization", "Bearer access-token-account-123")) + .and(header("chatgpt-account-id", "account-123")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "agent_runtime_id": "agent_123", + }))) + .expect(1) + .mount(&server) + .await; + + let tempdir = tempfile::tempdir().expect("tempdir"); + let keyring_store = Arc::new(MockKeyringStore::default()); + let secrets_manager = SecretsManager::new_with_keyring_store( + tempdir.path().to_path_buf(), + SecretsBackendKind::Local, + keyring_store, + ); + 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, + format!("{}/backend-api/", server.uri()), + SessionSource::Cli, + secrets_manager, + ); + + 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; + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "agent_runtime_id": "agent_456", + }))) + .expect(1) + .mount(&server) + .await; + + let tempdir = tempfile::tempdir().expect("tempdir"); + let keyring_store = Arc::new(MockKeyringStore::default()); + let secrets_manager = SecretsManager::new_with_keyring_store( + tempdir.path().to_path_buf(), + SecretsBackendKind::Local, + keyring_store.clone(), + ); + 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, + format!("{}/backend-api/", server.uri()), + SessionSource::Cli, + secrets_manager.clone(), + ); + + let binding = AgentIdentityBinding::from_auth( + &make_chatgpt_auth("account-123", Some("user-123")), + None, + ) + .expect("binding"); + let scope = secret_scope(&binding).expect("scope"); + let name = agent_identity_secret_name().expect("secret name"); + secrets_manager + .set(&scope, &name, "{\"not\":\"valid\"}") + .expect("seed invalid secret"); + + let stored = manager + .ensure_registered_identity() + .await + .unwrap() + .expect("identity should be registered"); + + assert_eq!(stored.agent_runtime_id, "agent_456"); + let persisted = secrets_manager + .get(&scope, &name) + .expect("read secret") + .expect("secret"); + let parsed: StoredAgentIdentity = serde_json::from_str(&persisted).expect("json"); + assert_eq!(parsed.agent_runtime_id, "agent_456"); + } + + #[tokio::test] + async fn ensure_registered_identity_falls_back_to_backend_api_v1() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .respond_with(ResponseTemplate::new(404)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/backend-api/v1/agent/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "agent_runtime_id": "agent_fallback", + }))) + .expect(1) + .mount(&server) + .await; + + let tempdir = tempfile::tempdir().expect("tempdir"); + let keyring_store = Arc::new(MockKeyringStore::default()); + let secrets_manager = SecretsManager::new_with_keyring_store( + tempdir.path().to_path_buf(), + SecretsBackendKind::Local, + keyring_store, + ); + 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, + format!("{}/backend-api/", server.uri()), + SessionSource::Cli, + secrets_manager, + ); + + let stored = manager + .ensure_registered_identity() + .await + .unwrap() + .expect("identity should be registered"); + assert_eq!(stored.agent_runtime_id, "agent_fallback"); + } + + #[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()), + }; + 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 cd6f1e84fd..9846f275cc 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -14,6 +14,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; @@ -1477,6 +1478,23 @@ impl Session { }); } + fn start_agent_identity_registration(self: &Arc) { + let weak_sess = Arc::downgrade(self); + tokio::spawn(async move { + let Some(sess) = weak_sess.upgrade() else { + return; + }; + if let Err(error) = sess + .services + .agent_identity_manager + .ensure_registered_identity() + .await + { + warn!(error = %error, "agent identity registration failed"); + } + }); + } + #[allow(clippy::too_many_arguments)] fn make_turn_context( conversation_id: ThreadId, @@ -2019,6 +2037,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, @@ -2116,6 +2139,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 6e901a0346..669d99245f 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -2866,6 +2866,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, @@ -3711,6 +3716,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, diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 5e62199d2d..cd7a86d79b 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -5,6 +5,8 @@ // the TUI or the tracing stack). #![deny(clippy::print_stdout, clippy::print_stderr)] +mod account; +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 2e89585f99..4b6c5b30f5 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,