diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 50dcc90e19..326adea6e4 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -243,6 +243,7 @@ fn dummy_chatgpt_auth_does_not_create_cwd_auth_json_when_identity_is_set() { agent_runtime_id: "agent_123".to_string(), agent_private_key: "pkcs8-base64".to_string(), registered_at: "2026-04-13T12:00:00Z".to_string(), + background_task_id: None, }; auth.set_agent_identity(record.clone()) diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index 5b752ee672..9b9e587c93 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -1,4 +1,5 @@ use chrono::DateTime; +use chrono::SecondsFormat; use chrono::Utc; use serde::Deserialize; use serde::Serialize; @@ -19,6 +20,7 @@ use std::sync::Mutex; use tracing::warn; use crate::token_data::TokenData; +use crate::token_data::decode_jwt_payload; use codex_app_server_protocol::AuthMode; use codex_config::types::AuthCredentialsStoreMode; use codex_keyring_store::DefaultKeyringStore; @@ -40,7 +42,11 @@ pub struct AuthDotJson { #[serde(default, skip_serializing_if = "Option::is_none")] pub last_refresh: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "deserialize_agent_identity" + )] pub agent_identity: Option, } @@ -56,6 +62,68 @@ pub struct AgentIdentityAuthRecord { pub background_task_id: Option, } +#[derive(Deserialize)] +#[serde(untagged)] +enum AgentIdentityAuthInput { + Record(AgentIdentityAuthRecord), + Jwt(String), +} + +#[derive(Deserialize)] +struct AgentIdentityJwtClaims { + workspace_id: String, + #[serde(default)] + chatgpt_user_id: Option, + agent_runtime_id: String, + agent_private_key: String, + #[serde(default)] + registered_at: Option, +} + +impl TryFrom for AgentIdentityAuthRecord { + type Error = std::io::Error; + + fn try_from(input: AgentIdentityAuthInput) -> std::io::Result { + match input { + AgentIdentityAuthInput::Record(record) => Ok(record), + AgentIdentityAuthInput::Jwt(jwt) => { + AgentIdentityAuthRecord::from_agent_identity_jwt(&jwt) + } + } + } +} + +impl AgentIdentityAuthRecord { + pub(crate) fn from_agent_identity_jwt(jwt: &str) -> std::io::Result { + let claims: AgentIdentityJwtClaims = + decode_jwt_payload(jwt).map_err(std::io::Error::other)?; + + Ok(Self { + workspace_id: claims.workspace_id, + chatgpt_user_id: claims.chatgpt_user_id, + agent_runtime_id: claims.agent_runtime_id, + agent_private_key: claims.agent_private_key, + registered_at: claims + .registered_at + .unwrap_or_else(|| Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true)), + background_task_id: None, + }) + } +} + +fn deserialize_agent_identity<'de, D>( + deserializer: D, +) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let input = Option::::deserialize(deserializer)?; + input + .map(AgentIdentityAuthRecord::try_from) + .transpose() + .map_err(serde::de::Error::custom) +} + pub(super) fn get_auth_file(codex_home: &Path) -> PathBuf { codex_home.join("auth.json") } diff --git a/codex-rs/login/src/auth/storage_tests.rs b/codex-rs/login/src/auth/storage_tests.rs index e00f09d464..aaa8fd9ecf 100644 --- a/codex-rs/login/src/auth/storage_tests.rs +++ b/codex-rs/login/src/auth/storage_tests.rs @@ -79,6 +79,64 @@ async fn file_storage_persists_agent_identity() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn file_storage_loads_agent_identity_from_jwt() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let expected_record = AgentIdentityAuthRecord { + workspace_id: "account-123".to_string(), + chatgpt_user_id: None, + agent_runtime_id: "agent_123".to_string(), + agent_private_key: "pkcs8-base64".to_string(), + registered_at: "2026-04-13T12:00:00Z".to_string(), + background_task_id: None, + }; + let agent_identity_jwt = jwt_with_payload(json!({ + "workspace_id": expected_record.workspace_id, + "agent_runtime_id": expected_record.agent_runtime_id, + "agent_private_key": expected_record.agent_private_key, + "registered_at": expected_record.registered_at, + })); + let auth_file = get_auth_file(codex_home.path()); + std::fs::write( + &auth_file, + serde_json::to_string_pretty(&json!({ + "auth_mode": "chatgpt", + "agent_identity": agent_identity_jwt, + }))?, + )?; + + let loaded = storage.load()?; + + assert_eq!( + loaded.expect("auth should load").agent_identity, + Some(expected_record) + ); + Ok(()) +} + +#[tokio::test] +async fn file_storage_rejects_invalid_agent_identity_jwt() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let auth_file = get_auth_file(codex_home.path()); + std::fs::write( + &auth_file, + serde_json::to_string_pretty(&json!({ + "auth_mode": "chatgpt", + "agent_identity": "not-a-jwt", + }))?, + )?; + + let err = storage.load().expect_err("invalid JWT should fail load"); + + assert!( + err.to_string().contains("invalid ID token format"), + "unexpected error: {err}" + ); + Ok(()) +} + #[test] fn file_storage_delete_removes_auth_file() -> anyhow::Result<()> { let dir = tempdir()?; @@ -199,6 +257,24 @@ fn id_token_with_prefix(prefix: &str) -> IdTokenInfo { crate::token_data::parse_chatgpt_jwt_claims(&fake_jwt).expect("fake JWT should parse") } +fn jwt_with_payload(payload: serde_json::Value) -> String { + #[derive(Serialize)] + struct Header { + alg: &'static str, + typ: &'static str, + } + + let header = Header { + alg: "none", + typ: "JWT", + }; + let encode = |bytes: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes); + let header_b64 = encode(&serde_json::to_vec(&header).expect("serialize header")); + let payload_b64 = encode(&serde_json::to_vec(&payload).expect("serialize payload")); + let signature_b64 = encode(b"sig"); + format!("{header_b64}.{payload_b64}.{signature_b64}") +} + fn auth_with_prefix(prefix: &str) -> AuthDotJson { AuthDotJson { auth_mode: Some(AuthMode::ApiKey), diff --git a/codex-rs/login/src/token_data.rs b/codex-rs/login/src/token_data.rs index 2952ecb2a7..4833f43eed 100644 --- a/codex-rs/login/src/token_data.rs +++ b/codex-rs/login/src/token_data.rs @@ -114,7 +114,7 @@ pub enum IdTokenInfoError { Json(#[from] serde_json::Error), } -fn decode_jwt_payload(jwt: &str) -> Result { +pub(crate) fn decode_jwt_payload(jwt: &str) -> Result { // JWT format: header.payload.signature let mut parts = jwt.split('.'); let (_header_b64, payload_b64, _sig_b64) = match (parts.next(), parts.next(), parts.next()) {