Register agent tasks behind use_agent_identity

codex-rs/Cargo.toml

@@ -219,6 +219,7 @@ color-eyre = "0.6.3"
 constant_time_eq = "0.3.1"
 crossbeam-channel = "0.5.15"
 crossterm = "0.28.1"
+crypto_box = { version = "0.9.1", features = ["seal"] }
 csv = "1.3.1"
 ctor = "0.6.3"
 deno_core_icudata = "0.77.0"

codex-rs/core/Cargo.toml

@@ -34,6 +34,7 @@ codex-code-mode = { workspace = true }
 codex-connectors = { workspace = true }
 codex-config = { workspace = true }
 codex-core-skills = { workspace = true }
+crypto_box = { workspace = true }
 codex-exec-server = { workspace = true }
 codex-features = { workspace = true }
 codex-feedback = { workspace = true }
@@ -97,6 +98,7 @@ rmcp = { workspace = true, default-features = false, features = [
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true }
 sha1 = { workspace = true }
+sha2 = { workspace = true }
 shlex = { workspace = true }
 similar = { workspace = true }
 tempfile = { workspace = true }

codex-rs/core/src/agent_identity.rs

@@ -27,6 +27,10 @@ use tracing::debug;
 use tracing::info;
 use tracing::warn;
 
+mod task_registration;
+
+pub(crate) use task_registration::RegisteredAgentTask;
+
 use crate::config::Config;
 
 const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15);
@@ -119,32 +123,58 @@ impl AgentIdentityManager {
             return Ok(None);
         }
 
-        let Some(auth) = self.auth_manager.auth().await else {
-            debug!("skipping agent identity registration because no auth is available");
+        let Some((auth, binding)) = self.current_auth_binding().await else {
             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);
-        };
+        self.ensure_registered_identity_for_binding(&auth, &binding)
+            .await
+            .map(Some)
+    }
 
+    async fn ensure_registered_identity_for_binding(
+        &self,
+        auth: &CodexAuth,
+        binding: &AgentIdentityBinding,
+    ) -> Result<StoredAgentIdentity> {
         let _guard = self.ensure_lock.lock().await;
 
-        if let Some(stored_identity) = self.load_stored_identity(&auth, &binding)? {
+        if let Some(stored_identity) = self.load_stored_identity(auth, binding)? {
             info!(
                 agent_runtime_id = %stored_identity.agent_runtime_id,
                 binding_id = %binding.binding_id,
                 "reusing stored agent identity"
             );
-            return Ok(Some(stored_identity));
+            return Ok(stored_identity);
         }
 
-        let stored_identity = self.register_agent_identity(&binding).await?;
-        self.store_identity(&auth, &stored_identity)?;
-        Ok(Some(stored_identity))
+        let stored_identity = self.register_agent_identity(binding).await?;
+        self.store_identity(auth, &stored_identity)?;
+        Ok(stored_identity)
+    }
+
+    pub(crate) async fn task_matches_current_binding(&self, task: &RegisteredAgentTask) -> bool {
+        if !self.feature_enabled {
+            return false;
+        }
+
+        self.current_auth_binding()
+            .await
+            .is_some_and(|(_, binding)| task.matches_binding(&binding))
+    }
+
+    async fn current_auth_binding(&self) -> Option<(CodexAuth, AgentIdentityBinding)> {
+        let Some(auth) = self.auth_manager.auth().await else {
+            debug!("skipping agent identity flow because no auth is available");
+            return None;
+        };
+
+        let binding =
+            AgentIdentityBinding::from_auth(&auth, self.auth_manager.forced_chatgpt_workspace_id());
+        if binding.is_none() {
+            debug!("skipping agent identity flow because ChatGPT auth is unavailable");
+        }
+        binding.map(|binding| (auth, binding))
     }
 
     async fn register_agent_identity(
@@ -351,12 +381,11 @@ 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,
-            }
+        binding.matches_parts(
+            &self.binding_id,
+            &self.chatgpt_account_id,
+            self.chatgpt_user_id.as_deref(),
+        )
     }
 
     fn validate_key_material(&self) -> Result<()> {
@@ -375,6 +404,20 @@ impl StoredAgentIdentity {
 }
 
 impl AgentIdentityBinding {
+    fn matches_parts(
+        &self,
+        binding_id: &str,
+        chatgpt_account_id: &str,
+        chatgpt_user_id: Option<&str>,
+    ) -> bool {
+        binding_id == self.binding_id
+            && chatgpt_account_id == self.chatgpt_account_id
+            && match self.chatgpt_user_id.as_deref() {
+                Some(expected_user_id) => chatgpt_user_id == Some(expected_user_id),
+                None => true,
+            }
+    }
+
     fn from_auth(auth: &CodexAuth, forced_workspace_id: Option<String>) -> Option<Self> {
         if !auth.is_chatgpt_auth() {
             return None;