Revalidate cached agent identity tasks

codex-rs/core/src/agent_identity.rs

@@ -153,6 +153,16 @@ impl AgentIdentityManager {
         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");
@@ -369,12 +379,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<()> {
@@ -393,6 +402,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;

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,6 +44,15 @@ impl AgentIdentityManager {
         let Some((auth, binding)) = self.current_auth_binding().await else {
             return Ok(None);
         };
+
+        self.register_task_for_binding(auth, binding).await
+    }
+
+    async fn register_task_for_binding(
+        &self,
+        auth: CodexAuth,
+        binding: AgentIdentityBinding,
+    ) -> Result<Option<RegisteredAgentTask>> {
         let stored_identity = self
             .ensure_registered_identity_for_binding(&auth, &binding)
             .await?;
@@ -70,6 +82,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,
@@ -91,6 +106,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,
@@ -211,6 +242,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(),
@@ -256,6 +290,79 @@ 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;
+        let target_url = agent_task_registration_url(&server.uri(), "agent-123");
+        mount_human_biscuit(&server, &target_url).await;
+        let binding_auth = make_chatgpt_auth("account-123", Some("user-123"));
+        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,
+        );
+        let stored_identity =
+            seed_stored_identity(&manager, &binding_auth, "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(&binding_auth, /*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_auth, 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 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,
+        );
+        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, target_url: &str) {
         Mock::given(method("GET"))
             .and(path("/authenticate_app_v2"))

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 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());
+            }
+
+            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));
         }
 
-        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));
-            }
-            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"
-        );
-        Ok(Some(agent_task))
+        Ok(None)
     }
 
     #[allow(clippy::too_many_arguments)]

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

codex-rs/core/src/state/session_tests.rs

@@ -39,6 +39,9 @@ async fn set_agent_task_persists_plaintext_task_for_session_reuse() {
     let session_configuration = make_session_configuration_for_tests().await;
     let mut state = SessionState::new(session_configuration);
     let agent_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(),
@@ -49,6 +52,25 @@ async fn set_agent_task_persists_plaintext_task_for_session_reuse() {
     assert_eq!(state.agent_task(), Some(agent_task));
 }
 
+#[tokio::test]
+async fn clear_agent_task_removes_cached_task() {
+    let session_configuration = make_session_configuration_for_tests().await;
+    let mut state = SessionState::new(session_configuration);
+    let agent_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(),
+    };
+
+    state.set_agent_task(agent_task);
+    state.clear_agent_task();
+
+    assert_eq!(state.agent_task(), None);
+}
+
 #[tokio::test]
 async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() {
     let session_configuration = make_session_configuration_for_tests().await;