diff --git a/codex-rs/agent-identity/src/lib.rs b/codex-rs/agent-identity/src/lib.rs index 2ce2abd495..243cec2a0c 100644 --- a/codex-rs/agent-identity/src/lib.rs +++ b/codex-rs/agent-identity/src/lib.rs @@ -234,14 +234,14 @@ struct RegisterTaskResponse { } #[derive(Debug, Serialize)] -struct RegisterAgentRequest { - abom: AgentBillOfMaterials, +struct CreateCodexAgentIdentityRequest { agent_public_key: String, - capabilities: Vec, + ttl: Option, + name: String, } #[derive(Debug, Deserialize)] -struct RegisterAgentResponse { +struct CreateCodexAgentIdentityResponse { agent_runtime_id: String, } @@ -398,13 +398,13 @@ pub async fn register_agent_identity( account_id: &str, is_fedramp_account: bool, key_material: &GeneratedAgentKeyMaterial, - abom: AgentBillOfMaterials, + name: &str, ) -> Result { let url = agent_registration_url(chatgpt_base_url); - let request = RegisterAgentRequest { - abom, + let request = CreateCodexAgentIdentityRequest { agent_public_key: key_material.public_key_ssh.clone(), - capabilities: Vec::new(), + ttl: None, + name: name.to_string(), }; let mut request_builder = client @@ -423,7 +423,7 @@ pub async fn register_agent_identity( .with_context(|| format!("failed to send agent identity registration request to {url}"))? .error_for_status() .with_context(|| format!("agent identity registration failed for {url}"))? - .json::() + .json::() .await .with_context(|| format!("failed to parse agent identity response from {url}"))?; @@ -496,21 +496,45 @@ pub fn curve25519_secret_key_from_private_key_pkcs8_base64( } pub fn agent_registration_url(chatgpt_base_url: &str) -> String { - let trimmed = chatgpt_base_url.trim_end_matches('/'); - format!("{trimmed}/v1/agent/register") + format!("{}/agent-identities", codex_api_base_url(chatgpt_base_url)) } pub fn agent_task_registration_url(chatgpt_base_url: &str, agent_runtime_id: &str) -> String { let trimmed = chatgpt_base_url.trim_end_matches('/'); - format!("{trimmed}/v1/agent/{agent_runtime_id}/task/register") + let api_path = format!("/v1/agent/{agent_runtime_id}/task/register"); + if matches!( + trimmed, + "https://chatgpt.com/backend-api" | "https://chat.openai.com/backend-api" + ) { + return format!("https://auth.openai.com/api/accounts{api_path}"); + } + if trimmed == "https://chatgpt-staging.com/backend-api" { + return format!("https://auth.api.openai.org/api/accounts{api_path}"); + } + if matches!( + trimmed, + "https://auth.openai.com" | "https://auth.api.openai.org" + ) { + return format!("{trimmed}/api/accounts{api_path}"); + } + format!("{trimmed}{api_path}") } pub fn agent_identity_jwks_url(chatgpt_base_url: &str) -> String { + format!( + "{}/agent-identities/jwks", + codex_api_base_url(chatgpt_base_url) + ) +} + +fn codex_api_base_url(chatgpt_base_url: &str) -> String { let trimmed = chatgpt_base_url.trim_end_matches('/'); - if trimmed.contains("/backend-api") { - format!("{trimmed}/wham/agent-identities/jwks") + if trimmed.ends_with("/api/codex") || trimmed.ends_with("/backend-api/codex") { + trimmed.to_string() + } else if trimmed.ends_with("/backend-api") { + format!("{trimmed}/codex") } else { - format!("{trimmed}/agent-identities/jwks") + format!("{trimmed}/api/codex") } } @@ -1009,15 +1033,46 @@ J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN ); } + #[test] + fn agent_registration_url_uses_codex_backend_api_base_url() { + assert_eq!( + agent_registration_url("https://chatgpt.com/backend-api"), + "https://chatgpt.com/backend-api/codex/agent-identities" + ); + assert_eq!( + agent_registration_url("http://localhost:8080/api/codex/"), + "http://localhost:8080/api/codex/agent-identities" + ); + assert_eq!( + agent_registration_url("http://localhost:8080"), + "http://localhost:8080/api/codex/agent-identities" + ); + } + + #[test] + fn agent_task_registration_url_uses_public_authapi_for_chatgpt_base_url() { + assert_eq!( + agent_task_registration_url("https://chatgpt.com/backend-api", "agent-runtime-id"), + "https://auth.openai.com/api/accounts/v1/agent/agent-runtime-id/task/register" + ); + assert_eq!( + agent_task_registration_url( + "https://chatgpt-staging.com/backend-api", + "agent-runtime-id" + ), + "https://auth.api.openai.org/api/accounts/v1/agent/agent-runtime-id/task/register" + ); + } + #[test] fn agent_identity_jwks_url_uses_backend_api_base_url() { assert_eq!( agent_identity_jwks_url("https://chatgpt.com/backend-api"), - "https://chatgpt.com/backend-api/wham/agent-identities/jwks" + "https://chatgpt.com/backend-api/codex/agent-identities/jwks" ); assert_eq!( agent_identity_jwks_url("https://chatgpt.com/backend-api/"), - "https://chatgpt.com/backend-api/wham/agent-identities/jwks" + "https://chatgpt.com/backend-api/codex/agent-identities/jwks" ); } diff --git a/codex-rs/cloud-requirements/src/lib.rs b/codex-rs/cloud-requirements/src/lib.rs index 8812762cbe..c05cc449ce 100644 --- a/codex-rs/cloud-requirements/src/lib.rs +++ b/codex-rs/cloud-requirements/src/lib.rs @@ -838,34 +838,13 @@ mod tests { use serde_json::json; use std::collections::BTreeMap; use std::collections::VecDeque; - use std::ffi::OsString; use std::future::pending; - use std::io::Read; - use std::io::Write; - use std::net::TcpListener; use std::path::Path; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; - use std::thread; use tempfile::TempDir; use tempfile::tempdir; - struct EnvVarGuard { - key: &'static str, - original: Option, - } - - impl Drop for EnvVarGuard { - fn drop(&mut self) { - unsafe { - match &self.original { - Some(value) => std::env::set_var(self.key, value), - None => std::env::remove_var(self.key), - } - } - } - } - fn write_auth_json(codex_home: &Path, value: serde_json::Value) -> std::io::Result<()> { std::fs::write(codex_home.join("auth.json"), serde_json::to_string(&value)?)?; Ok(()) @@ -1221,25 +1200,6 @@ mod tests { #[tokio::test] async fn cloud_requirements_eligible_auth_allows_agent_identity_business_plan() { - let listener = TcpListener::bind("127.0.0.1:0").expect("bind task registration server"); - let addr = listener - .local_addr() - .expect("task registration server addr"); - let server = thread::spawn(move || { - let (mut stream, _) = listener.accept().expect("accept task registration request"); - let mut request = [0; 4096]; - let _ = stream - .read(&mut request) - .expect("read task registration request"); - let body = r#"{"task_id":"task-123"}"#; - write!( - stream, - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - body.len(), - body - ) - .expect("write task registration response"); - }); let record = AgentIdentityAuthRecord { agent_runtime_id: "agent-runtime-123".to_string(), agent_private_key: "MC4CAQAwBQYDK2VwBCIEIDQg14jybCLydjHQwXeBzsDM7oB6BSAenodx6oCovQ/D" @@ -1249,21 +1209,9 @@ mod tests { email: "user@example.com".to_string(), plan_type: PlanType::Business, chatgpt_account_is_fedramp: false, + registered_at: None, }; - let authapi_base_url = format!("http://{addr}/backend-api"); - let original_authapi_base_url = std::env::var_os("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL"); - unsafe { - std::env::set_var("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &authapi_base_url); - } - let _authapi_guard = EnvVarGuard { - key: "CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", - original: original_authapi_base_url, - }; - let auth = AgentIdentityAuth::load(record) - .await - .map(CodexAuth::AgentIdentity) - .expect("agent identity auth"); - server.join().expect("task registration server joined"); + let auth = CodexAuth::AgentIdentity(AgentIdentityAuth::new(record)); assert!(cloud_requirements_eligible_auth(&auth)); } diff --git a/codex-rs/login/src/auth/agent_identity.rs b/codex-rs/login/src/auth/agent_identity.rs index 3644713328..3f55aaf901 100644 --- a/codex-rs/login/src/auth/agent_identity.rs +++ b/codex-rs/login/src/auth/agent_identity.rs @@ -1,43 +1,62 @@ +use std::sync::Arc; + use codex_agent_identity::AgentIdentityKey; +use codex_agent_identity::normalize_chatgpt_base_url; use codex_agent_identity::register_agent_task; use codex_protocol::account::PlanType as AccountPlanType; -use std::env; +use tokio::sync::OnceCell; use crate::default_client::build_reqwest_client; use super::storage::AgentIdentityAuthRecord; -const PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts"; -const CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR: &str = "CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL"; +const DEFAULT_CHATGPT_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api"; -#[derive(Clone, Debug)] +#[derive(Debug)] pub struct AgentIdentityAuth { record: AgentIdentityAuthRecord, - process_task_id: String, + process_task_id: Arc>, +} + +impl Clone for AgentIdentityAuth { + fn clone(&self) -> Self { + Self { + record: self.record.clone(), + process_task_id: Arc::clone(&self.process_task_id), + } + } } impl AgentIdentityAuth { - pub async fn load(record: AgentIdentityAuthRecord) -> std::io::Result { - let agent_identity_authapi_base_url = agent_identity_authapi_base_url(); - let process_task_id = register_agent_task( - &build_reqwest_client(), - &agent_identity_authapi_base_url, - key(&record), - ) - .await - .map_err(std::io::Error::other)?; - Ok(Self { + pub fn new(record: AgentIdentityAuthRecord) -> Self { + Self { record, - process_task_id, - }) + process_task_id: Arc::new(OnceCell::new()), + } } pub fn record(&self) -> &AgentIdentityAuthRecord { &self.record } - pub fn process_task_id(&self) -> &str { - &self.process_task_id + pub fn process_task_id(&self) -> Option<&str> { + self.process_task_id.get().map(String::as_str) + } + + pub async fn ensure_runtime(&self, chatgpt_base_url: Option) -> std::io::Result<()> { + self.process_task_id + .get_or_try_init(|| async { + let base_url = normalize_chatgpt_base_url( + chatgpt_base_url + .as_deref() + .unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL), + ); + register_agent_task(&build_reqwest_client(), &base_url, self.key()) + .await + .map_err(std::io::Error::other) + }) + .await + .map(|_| ()) } pub fn account_id(&self) -> &str { @@ -59,82 +78,10 @@ impl AgentIdentityAuth { pub fn is_fedramp_account(&self) -> bool { self.record.chatgpt_account_is_fedramp } -} - -fn agent_identity_authapi_base_url() -> String { - env::var(CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR) - .ok() - .map(|base_url| base_url.trim().trim_end_matches('/').to_string()) - .filter(|base_url| !base_url.is_empty()) - .unwrap_or_else(|| PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL.to_string()) -} - -fn key(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> { - AgentIdentityKey { - agent_runtime_id: &record.agent_runtime_id, - private_key_pkcs8_base64: &record.agent_private_key, - } -} - -#[cfg(test)] -mod tests { - use super::*; - use serial_test::serial; - - #[test] - #[serial(codex_auth_env)] - fn agent_identity_authapi_base_url_prefers_env_value() { - let _guard = EnvVarGuard::set( - CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR, - "https://authapi.example.test/api/accounts/", - ); - assert_eq!( - agent_identity_authapi_base_url(), - "https://authapi.example.test/api/accounts" - ); - } - - #[test] - #[serial(codex_auth_env)] - fn agent_identity_authapi_base_url_uses_prod_authapi_by_default() { - let _guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR); - assert_eq!( - agent_identity_authapi_base_url(), - PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL - ); - } - - struct EnvVarGuard { - key: &'static str, - original: Option, - } - - impl EnvVarGuard { - fn set(key: &'static str, value: &str) -> Self { - let original = env::var_os(key); - unsafe { - env::set_var(key, value); - } - Self { key, original } - } - - fn remove(key: &'static str) -> Self { - let original = env::var_os(key); - unsafe { - env::remove_var(key); - } - Self { key, original } - } - } - - impl Drop for EnvVarGuard { - fn drop(&mut self) { - unsafe { - match &self.original { - Some(value) => env::set_var(self.key, value), - None => env::remove_var(self.key), - } - } + fn key(&self) -> AgentIdentityKey<'_> { + AgentIdentityKey { + agent_runtime_id: &self.record.agent_runtime_id, + private_key_pkcs8_base64: &self.record.agent_private_key, } } } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index fe57be06fa..3ad0b02469 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::auth::storage::AgentIdentityStorage; use crate::auth::storage::FileAuthStorage; use crate::auth::storage::get_auth_file; use crate::token_data::IdTokenInfo; @@ -6,6 +7,7 @@ use codex_app_server_protocol::AuthMode; use codex_protocol::account::PlanType as AccountPlanType; use codex_protocol::auth::KnownPlan as InternalKnownPlan; use codex_protocol::auth::PlanType as InternalPlanType; +use codex_protocol::protocol::SessionSource; use base64::Engine; use codex_protocol::config_types::ForcedLoginMethod; @@ -19,6 +21,8 @@ use tempfile::tempdir; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; +use wiremock::matchers::body_partial_json; +use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; @@ -92,7 +96,7 @@ async fn login_with_access_token_writes_only_token() { signed_agent_identity_jwt(&record, json!(record.plan_type)).expect("signed agent identity"); let server = MockServer::start().await; Mock::given(method("GET")) - .and(path("/backend-api/wham/agent-identities/jwks")) + .and(path("/backend-api/codex/agent-identities/jwks")) .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body())) .expect(1) .mount(&server) @@ -114,7 +118,9 @@ async fn login_with_access_token_writes_only_token() { .expect("auth.json should parse"); assert_eq!(auth.auth_mode, Some(AuthMode::AgentIdentity)); assert_eq!( - auth.agent_identity.as_deref(), + auth.agent_identity + .as_ref() + .and_then(AgentIdentityStorage::as_jwt), Some(agent_identity.as_str()) ); assert!(auth.tokens.is_none(), "tokens should be cleared"); @@ -142,6 +148,101 @@ async fn login_with_access_token_rejects_invalid_jwt() { ); } +#[tokio::test] +async fn chatgpt_auth_registers_agent_identity_when_enabled() -> anyhow::Result<()> { + let codex_home = tempdir()?; + 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(), + )?; + let auth = super::load_auth( + codex_home.path(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await? + .expect("auth should load"); + + assert!( + auth.agent_identity_auth( + AgentIdentityAuthPolicy::JwtOnly, + /*chatgpt_base_url*/ None, + /*forced_chatgpt_workspace_id*/ None, + SessionSource::Cli, + ) + .await? + .is_none() + ); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/codex/agent-identities")) + .and(header("authorization", "Bearer test-access-token")) + .and(header("chatgpt-account-id", "account-123")) + .and(body_partial_json(json!({ + "name": "Codex CLI", + "ttl": null, + }))) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "agent_runtime_id": "agent-runtime-123", + }))) + .expect(/*r*/ 1) + .mount(&server) + .await; + + let agent_auth = auth + .agent_identity_auth( + AgentIdentityAuthPolicy::JwtOrChatgpt, + Some(server.uri()), + /*forced_chatgpt_workspace_id*/ None, + SessionSource::Cli, + ) + .await? + .expect("agent identity should register"); + let reused = auth + .agent_identity_auth( + AgentIdentityAuthPolicy::JwtOrChatgpt, + Some(server.uri()), + /*forced_chatgpt_workspace_id*/ None, + SessionSource::Cli, + ) + .await? + .expect("agent identity should be reused"); + Mock::given(method("POST")) + .and(path("/v1/agent/agent-runtime-123/task/register")) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "task_id": "task-123", + }))) + .expect(/*r*/ 1) + .mount(&server) + .await; + + agent_auth.ensure_runtime(Some(server.uri())).await?; + reused.ensure_runtime(Some(server.uri())).await?; + + assert_eq!( + agent_auth.record().agent_runtime_id, + reused.record().agent_runtime_id + ); + assert_eq!(agent_auth.process_task_id(), Some("task-123")); + assert_eq!(reused.process_task_id(), Some("task-123")); + assert_eq!(agent_auth.record().agent_runtime_id, "agent-runtime-123"); + assert_eq!(agent_auth.record().account_id, "account-123"); + assert_eq!(agent_auth.record().chatgpt_user_id, "user-12345"); + assert_eq!( + auth.get_agent_identity("account-123") + .expect("identity should persist") + .agent_runtime_id, + "agent-runtime-123" + ); + Ok(()) +} + #[tokio::test] async fn login_with_access_token_rejects_unsigned_jwt() { let dir = tempdir().unwrap(); @@ -149,7 +250,7 @@ async fn login_with_access_token_rejects_unsigned_jwt() { let agent_identity = fake_agent_identity_jwt(&record).expect("fake agent identity"); let server = MockServer::start().await; Mock::given(method("GET")) - .and(path("/backend-api/wham/agent-identities/jwks")) + .and(path("/backend-api/codex/agent-identities/jwks")) .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body())) .expect(1) .mount(&server) @@ -718,7 +819,7 @@ async fn load_auth_reads_access_token_from_env() { .expect("signed agent identity"); let server = MockServer::start().await; Mock::given(method("GET")) - .and(path("/backend-api/wham/agent-identities/jwks")) + .and(path("/backend-api/codex/agent-identities/jwks")) .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body())) .expect(1) .mount(&server) @@ -734,8 +835,6 @@ async fn load_auth_reads_access_token_from_env() { let _access_token_guard = EnvVarGuard::set(CODEX_ACCESS_TOKEN_ENV_VAR, &agent_identity); let chatgpt_base_url = format!("{}/backend-api", server.uri()); - let _authapi_guard = - EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url); let auth = super::load_auth( codex_home.path(), /*enable_codex_api_key_env*/ false, @@ -750,7 +849,7 @@ async fn load_auth_reads_access_token_from_env() { panic!("env auth should load as agent identity"); }; assert_eq!(agent_identity.record(), &expected_record); - assert_eq!(agent_identity.process_task_id(), "task-123"); + assert_eq!(agent_identity.process_task_id(), Some("task-123")); assert!( !get_auth_file(codex_home.path()).exists(), "env auth should not write auth.json" @@ -927,6 +1026,7 @@ fn agent_identity_record(account_id: &str) -> AgentIdentityAuthRecord { email: "user@example.com".to_string(), plan_type: AccountPlanType::Pro, chatgpt_account_is_fedramp: false, + registered_at: None, } } @@ -1045,7 +1145,7 @@ async fn assert_agent_identity_plan_alias( let jwt = signed_agent_identity_jwt(&record, plan_type).expect("agent identity jwt"); let server = MockServer::start().await; Mock::given(method("GET")) - .and(path("/backend-api/wham/agent-identities/jwks")) + .and(path("/backend-api/codex/agent-identities/jwks")) .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body())) .expect(1) .mount(&server) @@ -1059,8 +1159,6 @@ async fn assert_agent_identity_plan_alias( .mount(&server) .await; let chatgpt_base_url = format!("{}/backend-api", server.uri()); - let _authapi_guard = - EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url); let auth = CodexAuth::from_agent_identity_jwt(&jwt, Some(&chatgpt_base_url)) .await .expect("agent identity auth"); diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 45bd55302b..9076299fd8 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -18,6 +18,10 @@ use tokio::sync::Semaphore; use codex_agent_identity::decode_agent_identity_jwt; use codex_agent_identity::fetch_agent_identity_jwks; +use codex_agent_identity::generate_agent_key_material; +use codex_agent_identity::normalize_chatgpt_base_url; +use codex_agent_identity::public_key_ssh_from_private_key_pkcs8_base64; +use codex_agent_identity::register_agent_identity; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::AuthMode as ApiAuthMode; use codex_protocol::config_types::ForcedLoginMethod; @@ -27,6 +31,7 @@ use super::external_bearer::BearerTokenRefresher; use super::revoke::revoke_auth_tokens; pub use crate::auth::agent_identity::AgentIdentityAuth; pub use crate::auth::storage::AgentIdentityAuthRecord; +use crate::auth::storage::AgentIdentityStorage; pub use crate::auth::storage::AuthDotJson; use crate::auth::storage::AuthStorageBackend; use crate::auth::storage::create_auth_storage; @@ -42,6 +47,7 @@ use codex_protocol::account::PlanType as AccountPlanType; use codex_protocol::auth::PlanType as InternalPlanType; use codex_protocol::auth::RefreshTokenFailedError; use codex_protocol::auth::RefreshTokenFailedReason; +use codex_protocol::protocol::SessionSource; use serde_json::Value; use thiserror::Error; @@ -54,6 +60,15 @@ pub enum CodexAuth { AgentIdentity(AgentIdentityAuth), } +/// Policy for resolving Agent Identity auth from a broader Codex auth snapshot. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AgentIdentityAuthPolicy { + /// Only use an existing Agent Identity JWT/runtime auth. + JwtOnly, + /// Use an Agent Identity JWT/runtime auth, or register one from ChatGPT auth. + JwtOrChatgpt, +} + impl PartialEq for CodexAuth { fn eq(&self, other: &Self) -> bool { self.api_auth_mode() == other.api_auth_mode() @@ -79,10 +94,38 @@ pub struct ChatgptAuthTokens { #[derive(Debug, Clone)] struct ChatgptAuthState { auth_dot_json: Arc>>, + agent_identity_auth: Arc>>, client: CodexHttpClient, } +impl ChatgptAuthState { + fn new(auth_dot_json: AuthDotJson) -> Self { + let agent_identity_auth = auth_dot_json + .agent_identity + .as_ref() + .and_then(AgentIdentityStorage::as_record) + .cloned() + .map(AgentIdentityAuth::new); + Self { + auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), + agent_identity_auth: Arc::new(Mutex::new(agent_identity_auth)), + client: create_client(), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct ChatgptAgentIdentityBinding { + account_id: String, + chatgpt_user_id: String, + email: String, + plan_type: AccountPlanType, + chatgpt_account_is_fedramp: bool, + access_token: String, +} + const TOKEN_REFRESH_INTERVAL: i64 = 8; +const DEFAULT_CHATGPT_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api"; const REFRESH_TOKEN_EXPIRED_MESSAGE: &str = "Your access token could not be refreshed because your refresh token has expired. Please log out and sign in again."; const REFRESH_TOKEN_REUSED_MESSAGE: &str = "Your access token could not be refreshed because your refresh token was already used. Please log out and sign in again."; @@ -90,7 +133,6 @@ const REFRESH_TOKEN_INVALIDATED_MESSAGE: &str = "Your access token could not be const REFRESH_TOKEN_UNKNOWN_MESSAGE: &str = "Your access token could not be refreshed. Please log out and sign in again."; const REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE: &str = "Your access token could not be refreshed because you have since logged out or signed in to another account. Please sign in again."; -const DEFAULT_CHATGPT_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api"; const REFRESH_TOKEN_URL: &str = "https://auth.openai.com/oauth/token"; pub(super) const REVOKE_TOKEN_URL: &str = "https://auth.openai.com/oauth/revoke"; pub const REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR: &str = "CODEX_REFRESH_TOKEN_URL_OVERRIDE"; @@ -203,7 +245,6 @@ impl CodexAuth { chatgpt_base_url: Option<&str>, ) -> std::io::Result { let auth_mode = auth_dot_json.resolved_mode(); - let client = create_client(); if auth_mode == ApiAuthMode::ApiKey { let Some(api_key) = auth_dot_json.openai_api_key.as_deref() else { return Err(std::io::Error::other("API key auth is missing a key.")); @@ -211,19 +252,20 @@ impl CodexAuth { return Ok(Self::from_api_key(api_key)); } if auth_mode == ApiAuthMode::AgentIdentity { - let Some(agent_identity) = auth_dot_json.agent_identity else { + let Some(agent_identity) = auth_dot_json + .agent_identity + .as_ref() + .and_then(AgentIdentityStorage::as_jwt) + else { return Err(std::io::Error::other( "agent identity auth is missing an agent identity token.", )); }; - return Self::from_agent_identity_jwt(&agent_identity, chatgpt_base_url).await; + return Self::from_agent_identity_jwt(agent_identity, chatgpt_base_url).await; } let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); - let state = ChatgptAuthState { - auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), - client, - }; + let state = ChatgptAuthState::new(auth_dot_json); match auth_mode { ApiAuthMode::Chatgpt => { @@ -261,7 +303,9 @@ impl CodexAuth { .trim_end_matches('/') .to_string(); let record = verified_agent_identity_record(jwt, &base_url).await?; - Ok(Self::AgentIdentity(AgentIdentityAuth::load(record).await?)) + let auth = AgentIdentityAuth::new(record); + auth.ensure_runtime(Some(base_url)).await?; + Ok(Self::AgentIdentity(auth)) } pub fn auth_mode(&self) -> AuthMode { @@ -408,6 +452,196 @@ impl CodexAuth { self.get_current_auth_json().and_then(|t| t.tokens) } + pub fn get_agent_identity(&self, account_id: &str) -> Option { + self.get_current_auth_json() + .and_then(|auth| auth.agent_identity) + .and_then(|identity| identity.as_record().cloned()) + .filter(|identity| identity.account_id == account_id) + } + + pub fn set_agent_identity(&self, record: AgentIdentityAuthRecord) -> std::io::Result<()> { + let agent_identity_auth = self.agent_identity_auth_for_record(record.clone())?; + match self { + Self::Chatgpt(auth) => auth + .update_auth_json(|auth_dot_json| { + auth_dot_json.agent_identity = Some(AgentIdentityStorage::Record(record)); + true + }) + .map(|_| ()), + Self::ChatgptAuthTokens(auth) => auth.update_auth_json_in_memory(|auth_dot_json| { + auth_dot_json.agent_identity = Some(AgentIdentityStorage::Record(record)); + }), + Self::ApiKey(_) | Self::AgentIdentity(_) => Ok(()), + }?; + self.set_cached_agent_identity_auth(Some(agent_identity_auth)) + } + + pub fn remove_agent_identity(&self) -> std::io::Result { + let removed = match self { + Self::Chatgpt(auth) => { + auth.update_auth_json(|auth_dot_json| auth_dot_json.agent_identity.take().is_some()) + } + Self::ChatgptAuthTokens(auth) => { + let mut removed = false; + auth.update_auth_json_in_memory(|auth_dot_json| { + removed = auth_dot_json.agent_identity.take().is_some(); + })?; + Ok(removed) + } + Self::ApiKey(_) | Self::AgentIdentity(_) => Ok(false), + }?; + if removed { + self.set_cached_agent_identity_auth(/*auth*/ None)?; + } + Ok(removed) + } + + fn cached_agent_identity_auth( + &self, + binding: &ChatgptAgentIdentityBinding, + ) -> Option { + let auth = self.cached_agent_identity_auth_value()?; + if agent_identity_record_matches_binding(auth.record(), binding) + && public_key_ssh_from_private_key_pkcs8_base64(&auth.record().agent_private_key) + .is_ok() + { + Some(auth) + } else { + None + } + } + + fn agent_identity_auth_for_record( + &self, + record: AgentIdentityAuthRecord, + ) -> std::io::Result { + if let Some(auth) = self.cached_agent_identity_auth_value() + && auth.record() == &record + { + return Ok(auth); + } + Ok(AgentIdentityAuth::new(record)) + } + + fn cached_agent_identity_auth_value(&self) -> Option { + let state = self.chatgpt_state()?; + let auth = state.agent_identity_auth.lock().ok()?; + auth.clone() + } + + fn set_cached_agent_identity_auth( + &self, + auth: Option, + ) -> std::io::Result<()> { + let Some(state) = self.chatgpt_state() else { + return Ok(()); + }; + let mut cached = state + .agent_identity_auth + .lock() + .map_err(|_| std::io::Error::other("failed to lock agent identity cache"))?; + *cached = auth; + Ok(()) + } + + fn chatgpt_state(&self) -> Option<&ChatgptAuthState> { + match self { + Self::Chatgpt(auth) => Some(&auth.state), + Self::ChatgptAuthTokens(auth) => Some(&auth.state), + Self::ApiKey(_) | Self::AgentIdentity(_) => None, + } + } + + pub async fn agent_identity_auth( + &self, + policy: AgentIdentityAuthPolicy, + chatgpt_base_url: Option, + forced_chatgpt_workspace_id: Option, + session_source: SessionSource, + ) -> std::io::Result> { + match self { + Self::AgentIdentity(auth) => Ok(Some(auth.clone())), + Self::ApiKey(_) => Ok(None), + Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => { + if policy == AgentIdentityAuthPolicy::JwtOnly { + return Ok(None); + } + self.ensure_chatgpt_agent_identity( + chatgpt_base_url, + forced_chatgpt_workspace_id, + session_source, + ) + .await + .map(Some) + } + } + } + + async fn ensure_chatgpt_agent_identity( + &self, + chatgpt_base_url: Option, + forced_chatgpt_workspace_id: Option, + session_source: SessionSource, + ) -> std::io::Result { + let binding = ChatgptAgentIdentityBinding::from_auth(self, forced_chatgpt_workspace_id) + .ok_or_else(|| std::io::Error::other("ChatGPT auth is unavailable"))?; + + if let Some(auth) = self.cached_agent_identity_auth(&binding) { + return Ok(auth); + } + + if let Some(record) = self.get_agent_identity(&binding.account_id) + && agent_identity_record_matches_binding(&record, &binding) + && public_key_ssh_from_private_key_pkcs8_base64(&record.agent_private_key).is_ok() + { + let auth = self.agent_identity_auth_for_record(record)?; + self.set_cached_agent_identity_auth(Some(auth.clone()))?; + return Ok(auth); + } + + let key_material = generate_agent_key_material().map_err(std::io::Error::other)?; + let base_url = normalize_chatgpt_base_url( + chatgpt_base_url + .as_deref() + .unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL), + ); + let agent_name = match session_source { + SessionSource::VSCode => "Codex App", + SessionSource::Cli + | SessionSource::Exec + | SessionSource::Mcp + | SessionSource::Custom(_) + | SessionSource::Internal(_) + | SessionSource::SubAgent(_) + | SessionSource::Unknown => "Codex CLI", + }; + let runtime_id = register_agent_identity( + &build_reqwest_client(), + &base_url, + &binding.access_token, + &binding.account_id, + binding.chatgpt_account_is_fedramp, + &key_material, + agent_name, + ) + .await + .map_err(std::io::Error::other)?; + let record = AgentIdentityAuthRecord { + agent_runtime_id: runtime_id.into_string(), + agent_private_key: key_material.private_key_pkcs8_base64, + account_id: binding.account_id, + chatgpt_user_id: binding.chatgpt_user_id, + email: binding.email, + plan_type: binding.plan_type, + chatgpt_account_is_fedramp: binding.chatgpt_account_is_fedramp, + registered_at: Some( + Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, /*use_z*/ true), + ), + }; + self.set_agent_identity(record.clone())?; + self.agent_identity_auth_for_record(record) + } + /// Consider this private to integration tests. pub fn create_dummy_chatgpt_auth_for_testing() -> Self { let auth_dot_json = AuthDotJson { @@ -423,11 +657,7 @@ impl CodexAuth { agent_identity: None, }; - let client = create_client(); - let state = ChatgptAuthState { - auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), - client, - }; + let state = ChatgptAuthState::new(auth_dot_json); let dummy_auth_id = NEXT_DUMMY_AUTH_ID.fetch_add(1, Ordering::Relaxed); let storage = create_auth_storage( PathBuf::from(format!("dummy-chatgpt-auth-{dummy_auth_id}")), @@ -443,6 +673,44 @@ impl CodexAuth { } } +impl ChatgptAgentIdentityBinding { + 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 account_id = forced_workspace_id + .filter(|value| !value.is_empty()) + .or(token_data + .account_id + .clone() + .filter(|value| !value.is_empty())) + .or(token_data.id_token.chatgpt_account_id.clone())?; + let chatgpt_user_id = token_data + .id_token + .chatgpt_user_id + .clone() + .filter(|value| !value.is_empty())?; + + Some(Self { + account_id, + chatgpt_user_id, + email: token_data.id_token.email.clone().unwrap_or_default(), + plan_type: auth.account_plan_type().unwrap_or(AccountPlanType::Unknown), + chatgpt_account_is_fedramp: auth.is_fedramp_account(), + access_token: token_data.access_token, + }) + } +} + +fn agent_identity_record_matches_binding( + record: &AgentIdentityAuthRecord, + binding: &ChatgptAgentIdentityBinding, +) -> bool { + record.account_id == binding.account_id && record.chatgpt_user_id == binding.chatgpt_user_id +} + impl ChatgptAuth { fn current_auth_json(&self) -> Option { #[expect(clippy::unwrap_used)] @@ -460,6 +728,45 @@ impl ChatgptAuth { fn client(&self) -> &CodexHttpClient { &self.state.client } + + fn update_auth_json( + &self, + update: impl FnOnce(&mut AuthDotJson) -> bool, + ) -> std::io::Result { + let mut guard = self + .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"))?; + let changed = update(&mut auth); + if changed { + self.storage.save(&auth)?; + *guard = Some(auth); + } + Ok(changed) + } +} + +impl ChatgptAuthTokens { + fn update_auth_json_in_memory( + &self, + update: impl FnOnce(&mut AuthDotJson), + ) -> std::io::Result<()> { + let mut guard = self + .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"))?; + update(&mut auth); + *guard = Some(auth); + Ok(()) + } } pub const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; @@ -558,7 +865,7 @@ pub async fn login_with_access_token( openai_api_key: None, tokens: None, last_refresh: None, - agent_identity: Some(access_token.to_string()), + agent_identity: Some(AgentIdentityStorage::Jwt(access_token.to_string())), }; save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode) } @@ -1250,6 +1557,7 @@ pub struct AuthManager { forced_chatgpt_workspace_id: RwLock>, chatgpt_base_url: Option, refresh_lock: Semaphore, + agent_identity_lock: Semaphore, external_auth: RwLock>>, } @@ -1324,6 +1632,7 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), chatgpt_base_url, refresh_lock: Semaphore::new(/*permits*/ 1), + agent_identity_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), } } @@ -1343,6 +1652,7 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), chatgpt_base_url: None, refresh_lock: Semaphore::new(/*permits*/ 1), + agent_identity_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), }) } @@ -1361,6 +1671,7 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), chatgpt_base_url: None, refresh_lock: Semaphore::new(/*permits*/ 1), + agent_identity_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), }) } @@ -1377,6 +1688,7 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), chatgpt_base_url: None, refresh_lock: Semaphore::new(/*permits*/ 1), + agent_identity_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(Some( Arc::new(BearerTokenRefresher::new(config)) as Arc )), @@ -1416,6 +1728,38 @@ impl AuthManager { self.auth_cached() } + pub async fn agent_identity_auth( + &self, + policy: AgentIdentityAuthPolicy, + session_source: SessionSource, + ) -> std::io::Result> { + let Some(auth) = self.auth().await else { + return Ok(None); + }; + if policy == AgentIdentityAuthPolicy::JwtOrChatgpt && auth.is_chatgpt_auth() { + let _permit = self + .agent_identity_lock + .acquire() + .await + .map_err(std::io::Error::other)?; + return auth + .agent_identity_auth( + policy, + self.chatgpt_base_url.clone(), + self.forced_chatgpt_workspace_id(), + session_source, + ) + .await; + } + auth.agent_identity_auth( + policy, + self.chatgpt_base_url.clone(), + self.forced_chatgpt_workspace_id(), + session_source, + ) + .await + } + /// Force a reload of the auth information from auth.json. Returns /// whether the auth value changed. pub async fn reload(&self) -> bool { diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index 3a1c8ae6aa..a0859419ab 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -44,7 +44,30 @@ pub struct AuthDotJson { pub last_refresh: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] - pub agent_identity: Option, + pub agent_identity: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] +#[serde(untagged)] +pub enum AgentIdentityStorage { + Jwt(String), + Record(AgentIdentityAuthRecord), +} + +impl AgentIdentityStorage { + pub(crate) fn as_jwt(&self) -> Option<&str> { + match self { + Self::Jwt(jwt) => Some(jwt), + Self::Record(_) => None, + } + } + + pub(crate) fn as_record(&self) -> Option<&AgentIdentityAuthRecord> { + match self { + Self::Jwt(_) => None, + Self::Record(record) => Some(record), + } + } } #[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] @@ -56,6 +79,8 @@ pub struct AgentIdentityAuthRecord { pub email: String, pub plan_type: AccountPlanType, pub chatgpt_account_is_fedramp: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub registered_at: Option, } impl AgentIdentityAuthRecord { @@ -77,6 +102,7 @@ impl From for AgentIdentityAuthRecord { email: claims.email, plan_type: claims.plan_type.into(), chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp, + registered_at: None, } } } diff --git a/codex-rs/login/src/auth/storage_tests.rs b/codex-rs/login/src/auth/storage_tests.rs index b5646ef53e..f6df9f2d74 100644 --- a/codex-rs/login/src/auth/storage_tests.rs +++ b/codex-rs/login/src/auth/storage_tests.rs @@ -72,7 +72,7 @@ async fn file_storage_round_trips_agent_identity_auth() -> anyhow::Result<()> { openai_api_key: None, tokens: None, last_refresh: None, - agent_identity: Some(agent_identity), + agent_identity: Some(AgentIdentityStorage::Jwt(agent_identity)), }; storage.save(&auth_dot_json)?; @@ -107,7 +107,11 @@ async fn file_storage_loads_agent_identity_as_jwt() -> anyhow::Result<()> { let loaded = storage.load()?; assert_eq!( - loaded.expect("auth should load").agent_identity.as_deref(), + loaded + .expect("auth should load") + .agent_identity + .as_ref() + .and_then(AgentIdentityStorage::as_jwt), Some(agent_identity_jwt.as_str()) ); Ok(()) diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 990cf8b80e..53f9b0ca92 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -17,6 +17,7 @@ pub use server::ServerOptions; pub use server::ShutdownHandle; pub use server::run_login_server; +pub use auth::AgentIdentityAuthPolicy; pub use auth::AuthConfig; pub use auth::AuthDotJson; pub use auth::AuthManager; diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 3c7f4dbd0f..2edc052dad 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -21,6 +21,9 @@ struct AgentIdentityAuthProvider { impl AuthProvider for AgentIdentityAuthProvider { fn add_auth_headers(&self, headers: &mut HeaderMap) { let record = self.auth.record(); + let Some(task_id) = self.auth.process_task_id() else { + return; + }; let header_value = authorization_header_for_agent_task( AgentIdentityKey { agent_runtime_id: &record.agent_runtime_id, @@ -28,7 +31,7 @@ impl AuthProvider for AgentIdentityAuthProvider { }, AgentTaskAuthorizationTarget { agent_runtime_id: &record.agent_runtime_id, - task_id: self.auth.process_task_id(), + task_id, }, ) .map_err(std::io::Error::other);