Revalidate cached agent identity tasks

adrian-openai · adrian-openai · commit 3cbb52231099 · 2026-04-10T23:30:10.000-07:00
diff --git a/codex-rs/core/src/agent_identity.rs b/codex-rs/core/src/agent_identity.rs
@@ -150,9 +150,20 @@ impl AgentIdentityManager {
             return Ok(None);
         };
 
+        self.ensure_registered_identity_for_binding(&binding).await
+    }
+
+    async fn ensure_registered_identity_for_binding(
+        &self,
+        binding: &AgentIdentityBinding,
+    ) -> Result<Option<StoredAgentIdentity>> {
+        if !self.feature_enabled {
+            return Ok(None);
+        }
+
         let _guard = self.ensure_lock.lock().await;
 
-        if let Some(stored_identity) = self.load_stored_identity(&binding)? {
+        if let Some(stored_identity) = self.load_stored_identity(binding)? {
             info!(
                 agent_runtime_id = %stored_identity.agent_runtime_id,
                 binding_id = %binding.binding_id,
@@ -161,11 +172,21 @@ impl AgentIdentityManager {
             return Ok(Some(stored_identity));
         }
 
-        let stored_identity = self.register_agent_identity(&binding).await?;
-        self.store_identity(&binding, &stored_identity)?;
+        let stored_identity = self.register_agent_identity(binding).await?;
+        self.store_identity(binding, &stored_identity)?;
         Ok(Some(stored_identity))
     }
 
+    pub(crate) async fn task_matches_current_binding(&self, task: &RegisteredAgentTask) -> bool {
+        if !self.feature_enabled {
+            return false;
+        }
+
+        self.current_binding()
+            .await
+            .is_some_and(|binding| task.matches_binding(&binding))
+    }
+
     async fn current_binding(&self) -> Option<AgentIdentityBinding> {
         let Some(auth) = self.auth_manager.auth().await else {
             debug!("skipping agent identity flow because no auth is available");
@@ -360,12 +381,11 @@ impl AgentIdentityManager {
 
 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<()> {
@@ -388,6 +408,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;
diff --git a/codex-rs/core/src/agent_identity/task_registration.rs b/codex-rs/core/src/agent_identity/task_registration.rs
@@ -16,6 +16,9 @@ const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15);
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub(crate) struct RegisteredAgentTask {
+    pub(crate) binding_id: String,
+    pub(crate) chatgpt_account_id: String,
+    pub(crate) chatgpt_user_id: Option<String>,
     pub(crate) agent_runtime_id: String,
     pub(crate) task_id: String,
     pub(crate) registered_at: String,
@@ -41,7 +44,18 @@ impl AgentIdentityManager {
         let Some(binding) = self.current_binding().await else {
             return Ok(None);
         };
-        let Some(stored_identity) = self.ensure_registered_identity().await? else {
+
+        self.register_task_for_binding(binding).await
+    }
+
+    async fn register_task_for_binding(
+        &self,
+        binding: AgentIdentityBinding,
+    ) -> Result<Option<RegisteredAgentTask>> {
+        let Some(stored_identity) = self
+            .ensure_registered_identity_for_binding(&binding)
+            .await?
+        else {
             return Ok(None);
         };
 
@@ -72,6 +86,9 @@ impl AgentIdentityManager {
                 .await
                 .with_context(|| format!("failed to parse agent task response from {url}"))?;
             let registered_task = RegisteredAgentTask {
+                binding_id: stored_identity.binding_id.clone(),
+                chatgpt_account_id: stored_identity.chatgpt_account_id.clone(),
+                chatgpt_user_id: stored_identity.chatgpt_user_id.clone(),
                 agent_runtime_id: stored_identity.agent_runtime_id.clone(),
                 task_id: decrypt_task_id_response(
                     &stored_identity,
@@ -93,6 +110,22 @@ impl AgentIdentityManager {
     }
 }
 
+impl RegisteredAgentTask {
+    pub(super) fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool {
+        binding.matches_parts(
+            &self.binding_id,
+            &self.chatgpt_account_id,
+            self.chatgpt_user_id.as_deref(),
+        )
+    }
+
+    pub(crate) fn has_same_binding(&self, other: &Self) -> bool {
+        self.binding_id == other.binding_id
+            && self.chatgpt_account_id == other.chatgpt_account_id
+            && self.chatgpt_user_id == other.chatgpt_user_id
+    }
+}
+
 fn sign_task_registration_payload(
     stored_identity: &StoredAgentIdentity,
     timestamp: &str,
@@ -240,6 +273,9 @@ mod tests {
         assert_eq!(
             task,
             RegisteredAgentTask {
+                binding_id: "chatgpt-account-account-123".to_string(),
+                chatgpt_account_id: "account-123".to_string(),
+                chatgpt_user_id: Some("user-123".to_string()),
                 agent_runtime_id: "agent-123".to_string(),
                 task_id: "task_123".to_string(),
                 registered_at: task.registered_at.clone(),
@@ -292,6 +328,95 @@ mod tests {
         assert_eq!(task.task_id, "task_fallback");
     }
 
+    #[tokio::test]
+    async fn register_task_for_binding_keeps_one_auth_snapshot() {
+        let server = MockServer::start().await;
+        mount_human_biscuit(&server).await;
+        let tempdir = tempfile::tempdir().expect("tempdir");
+        let keyring_store = Arc::new(MockKeyringStore::default());
+        let secrets_manager = SecretsManager::new_with_keyring_store(
+            tempdir.path().to_path_buf(),
+            SecretsBackendKind::Local,
+            keyring_store,
+        );
+        let auth_manager =
+            AuthManager::from_auth_for_testing(make_chatgpt_auth("account-456", Some("user-456")));
+        let manager = AgentIdentityManager::new_for_tests(
+            auth_manager,
+            /*feature_enabled*/ true,
+            server.uri(),
+            SessionSource::Cli,
+            secrets_manager.clone(),
+        );
+        let stored_identity =
+            seed_stored_identity(&manager, &secrets_manager, "agent-123", "account-123");
+        let encrypted_task_id =
+            encrypt_task_id_for_identity(&stored_identity, "task_123").expect("task ciphertext");
+        let binding = AgentIdentityBinding::from_auth(
+            &make_chatgpt_auth("account-123", Some("user-123")),
+            /*forced_workspace_id*/ None,
+        )
+        .expect("binding");
+
+        Mock::given(method("POST"))
+            .and(path("/v1/agent/agent-123/task/register"))
+            .and(header("x-openai-authorization", "human-biscuit"))
+            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
+                "encrypted_task_id": encrypted_task_id,
+            })))
+            .expect(1)
+            .mount(&server)
+            .await;
+
+        let task = manager
+            .register_task_for_binding(binding)
+            .await
+            .unwrap()
+            .expect("task should be registered");
+
+        assert_eq!(
+            task,
+            RegisteredAgentTask {
+                binding_id: "chatgpt-account-account-123".to_string(),
+                chatgpt_account_id: "account-123".to_string(),
+                chatgpt_user_id: Some("user-123".to_string()),
+                agent_runtime_id: "agent-123".to_string(),
+                task_id: "task_123".to_string(),
+                registered_at: task.registered_at.clone(),
+            }
+        );
+    }
+
+    #[tokio::test]
+    async fn task_matches_current_binding_rejects_stale_auth_binding() {
+        let tempdir = tempfile::tempdir().expect("tempdir");
+        let keyring_store = Arc::new(MockKeyringStore::default());
+        let secrets_manager = SecretsManager::new_with_keyring_store(
+            tempdir.path().to_path_buf(),
+            SecretsBackendKind::Local,
+            keyring_store,
+        );
+        let auth_manager =
+            AuthManager::from_auth_for_testing(make_chatgpt_auth("account-456", Some("user-456")));
+        let manager = AgentIdentityManager::new_for_tests(
+            auth_manager,
+            /*feature_enabled*/ true,
+            "https://chatgpt.com/backend-api/".to_string(),
+            SessionSource::Cli,
+            secrets_manager,
+        );
+        let task = RegisteredAgentTask {
+            binding_id: "chatgpt-account-account-123".to_string(),
+            chatgpt_account_id: "account-123".to_string(),
+            chatgpt_user_id: Some("user-123".to_string()),
+            agent_runtime_id: "agent-123".to_string(),
+            task_id: "task_123".to_string(),
+            registered_at: "2026-03-23T12:00:00Z".to_string(),
+        };
+
+        assert!(!manager.task_matches_current_binding(&task).await);
+    }
+
     async fn mount_human_biscuit(server: &MockServer) {
         Mock::given(method("GET"))
             .and(path("/authenticate_app_v2"))
diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs
@@ -1528,37 +1528,89 @@ impl Session {
         .await;
         handlers::shutdown(self, self.next_internal_sub_id()).await;
     }
-    async fn ensure_agent_task_registered(&self) -> anyhow::Result<Option<RegisteredAgentTask>> {
-        {
+
+    async fn cached_agent_task_for_current_binding(&self) -> Option<RegisteredAgentTask> {
+        let agent_task = {
             let state = self.state.lock().await;
-            if let Some(agent_task) = state.agent_task() {
+            state.agent_task()
+        }?;
+
+        if self
+            .services
+            .agent_identity_manager
+            .task_matches_current_binding(&agent_task)
+            .await
+        {
+            debug!(
+                agent_runtime_id = %agent_task.agent_runtime_id,
+                task_id = %agent_task.task_id,
+                "reusing cached agent task"
+            );
+            return Some(agent_task);
+        }
+
+        debug!(
+            agent_runtime_id = %agent_task.agent_runtime_id,
+            task_id = %agent_task.task_id,
+            "discarding cached agent task because auth binding changed"
+        );
+        let mut state = self.state.lock().await;
+        if state.agent_task().as_ref() == Some(&agent_task) {
+            state.clear_agent_task();
+        }
+        None
+    }
+
+    async fn ensure_agent_task_registered(&self) -> anyhow::Result<Option<RegisteredAgentTask>> {
+        if let Some(agent_task) = self.cached_agent_task_for_current_binding().await {
+            return Ok(Some(agent_task));
+        }
+
+        for _ in 0..2 {
+            let Some(agent_task) = self.services.agent_identity_manager.register_task().await?
+            else {
+                return Ok(None);
+            };
+
+            if !self
+                .services
+                .agent_identity_manager
+                .task_matches_current_binding(&agent_task)
+                .await
+            {
                 debug!(
                     agent_runtime_id = %agent_task.agent_runtime_id,
                     task_id = %agent_task.task_id,
-                    "reusing cached agent task"
+                    "discarding newly registered agent task because auth binding changed"
                 );
-                return Ok(Some(agent_task));
+                continue;
             }
-        }
 
-        let Some(agent_task) = self.services.agent_identity_manager.register_task().await? else {
-            return Ok(None);
-        };
-        {
-            let mut state = self.state.lock().await;
-            if let Some(existing_agent_task) = state.agent_task() {
-                return Ok(Some(existing_agent_task));
+            {
+                let mut state = self.state.lock().await;
+                if let Some(existing_agent_task) = state.agent_task() {
+                    if existing_agent_task.has_same_binding(&agent_task) {
+                        return Ok(Some(existing_agent_task));
+                    }
+                    debug!(
+                        agent_runtime_id = %existing_agent_task.agent_runtime_id,
+                        task_id = %existing_agent_task.task_id,
+                        "replacing cached agent task because auth binding changed"
+                    );
+                }
+                state.set_agent_task(agent_task.clone());
             }
-            state.set_agent_task(agent_task.clone());
+
+            info!(
+                thread_id = %self.conversation_id,
+                agent_runtime_id = %agent_task.agent_runtime_id,
+                task_id = %agent_task.task_id,
+                "registered agent task for thread"
+            );
+            return Ok(Some(agent_task));
         }
 
-        info!(
-            thread_id = %self.conversation_id,
-            agent_runtime_id = %agent_task.agent_runtime_id,
-            task_id = %agent_task.task_id,
-            "registered agent task for thread"
-        );
-        Ok(Some(agent_task))
+        Ok(None)
     }
 
     #[allow(clippy::too_many_arguments)]
diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs
@@ -185,6 +185,10 @@ impl SessionState {
         self.agent_task = Some(agent_task);
     }
 
+    pub(crate) fn clear_agent_task(&mut self) {
+        self.agent_task = None;
+    }
+
     // Adds connector IDs to the active set and returns the merged selection.
     pub(crate) fn merge_connector_selection<I>(&mut self, connector_ids: I) -> HashSet<String>
     where
diff --git a/codex-rs/core/src/state/session_tests.rs b/codex-rs/core/src/state/session_tests.rs
