fix: accept JWT agent identity auth input

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())

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<DateTime<Utc>>,
 
-    #[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<AgentIdentityAuthRecord>,
 }
 
@@ -56,6 +62,68 @@ pub struct AgentIdentityAuthRecord {
     pub background_task_id: Option<String>,
 }
 
+#[derive(Deserialize)]
+#[serde(untagged)]
+enum AgentIdentityAuthInput {
+    Record(AgentIdentityAuthRecord),
+    Jwt(String),
+}
+
+#[derive(Deserialize)]
+struct AgentIdentityJwtClaims {
+    workspace_id: String,
+    #[serde(default)]
+    chatgpt_user_id: Option<String>,
+    agent_runtime_id: String,
+    agent_private_key: String,
+    #[serde(default)]
+    registered_at: Option<String>,
+}
+
+impl TryFrom<AgentIdentityAuthInput> for AgentIdentityAuthRecord {
+    type Error = std::io::Error;
+
+    fn try_from(input: AgentIdentityAuthInput) -> std::io::Result<Self> {
+        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<Self> {
+        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<Option<AgentIdentityAuthRecord>, D::Error>
+where
+    D: serde::Deserializer<'de>,
+{
+    let input = Option::<AgentIdentityAuthInput>::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")
 }

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),

codex-rs/login/src/token_data.rs

@@ -114,7 +114,7 @@ pub enum IdTokenInfoError {
     Json(#[from] serde_json::Error),
 }
 
-fn decode_jwt_payload<T: DeserializeOwned>(jwt: &str) -> Result<T, IdTokenInfoError> {
+pub(crate) fn decode_jwt_payload<T: DeserializeOwned>(jwt: &str) -> Result<T, IdTokenInfoError> {
     // JWT format: header.payload.signature
     let mut parts = jwt.split('.');
     let (_header_b64, payload_b64, _sig_b64) = match (parts.next(), parts.next(), parts.next()) {