Register agent identities behind use_agent_identity (#17386)

## Summary Stack PR 2 of 4 for feature-gated agent identity support. This PR adds agent identity registration behind `features.use_agent_identity`. It keeps the app-server protocol unchanged and starts registration after ChatGPT auth exists rather than requiring a client restart. ## Stack - PR1: https://github.com/openai/codex/pull/17385 - add `features.use_agent_identity` - PR2: https://github.com/openai/codex/pull/17386 - this PR - PR3: https://github.com/openai/codex/pull/17387 - register agent tasks when enabled - PR4: https://github.com/openai/codex/pull/17388 - use `AgentAssertion` downstream when enabled ## Validation Covered as part of the local stack validation pass: - `just fmt` - `cargo test -p codex-core --lib agent_identity` - `cargo test -p codex-core --lib agent_assertion` - `cargo test -p codex-core --lib websocket_agent_task` - `cargo test -p codex-api api_bridge` - `cargo build -p codex-cli --bin codex` ## Notes The full local app-server E2E path is still being debugged after PR creation. The current branch stack is directionally ready for review while that follow-up continues.
2026-11-04 00:24:43 +00:00 · 2026-04-15 10:08:27 -07:00
1 parent 1dead46c90
commit 8e784bba2f
@@ -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<String>,
    chatgpt_plan_type: Option<String>,
@@ -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<dyn AuthStorageBackend>,
 }

 #[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<AgentIdentityAuthRecord> {
+        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<bool> {
+        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<Option<String>>,
    refresh_lock: AsyncMutex<()>,
    external_auth: RwLock<Option<Arc<dyn ExternalAuth>>>,
+    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<Self> {
+        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<Self> {
+        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<Self> {
+        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<dyn ExternalAuth>
            )),
+            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<dyn ExternalAuth>) {
        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<String>) {
-        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,
@@ -39,6 +39,19 @@ 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")]
+    pub agent_identity: Option<AgentIdentityAuthRecord>,
+}
+
+#[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<String>,
+    pub agent_runtime_id: String,
+    pub agent_private_key: String,
+    pub registered_at: String,
 }

 pub(super) fn get_auth_file(codex_home: &Path) -> PathBuf {
@@ -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)?;