codex: address PR review feedback (#14710)

codex-rs/tui/src/app.rs

@@ -3911,12 +3911,19 @@ impl App {
             self.suppress_shutdown_complete = false;
             return;
         }
-        if let ThreadUpdate::ThreadRolledBack { num_turns } = update {
+        if let ThreadUpdate::ThreadRolledBack {
+            thread_id,
+            num_turns,
+        } = update
+        {
             self.handle_backtrack_event(&codex_protocol::protocol::EventMsg::ThreadRolledBack(
                 codex_protocol::protocol::ThreadRolledBackEvent { num_turns },
             ));
             self.chat_widget
-                .handle_thread_update(ThreadUpdate::ThreadRolledBack { num_turns });
+                .handle_thread_update(ThreadUpdate::ThreadRolledBack {
+                    thread_id,
+                    num_turns,
+                });
             return;
         }
         self.chat_widget.handle_thread_update(update.clone());
@@ -4652,6 +4659,53 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn routed_thread_rollback_stays_on_target_thread() -> Result<()> {
+        let mut app = make_test_app().await;
+        let primary_thread_id = ThreadId::new();
+        let secondary_thread_id = ThreadId::new();
+        app.primary_thread_id = Some(primary_thread_id);
+        app.thread_event_channels.insert(
+            primary_thread_id,
+            ThreadEventChannel::new(THREAD_EVENT_CHANNEL_CAPACITY),
+        );
+        app.thread_event_channels.insert(
+            secondary_thread_id,
+            ThreadEventChannel::new(THREAD_EVENT_CHANNEL_CAPACITY),
+        );
+
+        app.route_thread_update(ThreadUpdate::ThreadRolledBack {
+            thread_id: secondary_thread_id.to_string(),
+            num_turns: 1,
+        })
+        .await;
+
+        let secondary_store = Arc::clone(
+            &app.thread_event_channels
+                .get(&secondary_thread_id)
+                .expect("missing secondary thread channel")
+                .store,
+        );
+        let secondary_snapshot = secondary_store.lock().await.snapshot();
+        assert!(matches!(
+            secondary_snapshot.updates.as_slice(),
+            [ThreadUpdate::ThreadRolledBack {
+                thread_id,
+                num_turns: 1,
+            }] if thread_id == &secondary_thread_id.to_string()
+        ));
+
+        let primary_store = Arc::clone(
+            &app.thread_event_channels
+                .get(&primary_thread_id)
+                .expect("missing primary thread channel")
+                .store,
+        );
+        let primary_snapshot = primary_store.lock().await.snapshot();
+        assert!(primary_snapshot.updates.is_empty());
+        Ok(())
+    }
+
     #[tokio::test]
     async fn reset_thread_event_state_aborts_listener_tasks() {
         struct NotifyOnDrop(Option<tokio::sync::oneshot::Sender<()>>);

codex-rs/tui/src/app/app_server_adapter.rs

@@ -692,8 +692,11 @@ impl App {
                     "thread/rollback",
                 )
                 .await?;
-                self.route_thread_update(ThreadUpdate::ThreadRolledBack { num_turns })
-                    .await;
+                self.route_thread_update(ThreadUpdate::ThreadRolledBack {
+                    thread_id: thread_id.to_string(),
+                    num_turns,
+                })
+                .await;
             }
             Op::Review { review_request } => {
                 let response: ReviewStartResponse = send_request_with_response(
@@ -819,7 +822,7 @@ impl App {
         request_id
     }
 
-    async fn route_thread_update(&mut self, update: ThreadUpdate) {
+    pub(super) async fn route_thread_update(&mut self, update: ThreadUpdate) {
         let Some(thread_id_text) = update.thread_id() else {
             self.handle_thread_update_now(update);
             return;

codex-rs/tui/src/chatwidget.rs

@@ -5079,7 +5079,7 @@ impl ChatWidget {
             ThreadUpdate::ThreadStarted(_)
             | ThreadUpdate::ThreadStatusChanged(_)
             | ThreadUpdate::ThreadClosed(_) => {}
-            ThreadUpdate::ThreadRolledBack { num_turns } => {
+            ThreadUpdate::ThreadRolledBack { num_turns, .. } => {
                 self.last_copyable_output = None;
                 if from_replay {
                     self.app_event_tx

codex-rs/tui/src/thread_update.rs

@@ -99,6 +99,7 @@ pub(crate) enum ThreadUpdate {
         params: DynamicToolCallParams,
     },
     ThreadRolledBack {
+        thread_id: String,
         num_turns: u32,
     },
 }
@@ -150,7 +151,8 @@ impl ThreadUpdate {
             Self::PermissionsRequestApproval { params, .. } => Some(params.thread_id.clone()),
             Self::ToolRequestUserInput { params, .. } => Some(params.thread_id.clone()),
             Self::DynamicToolCall { params, .. } => Some(params.thread_id.clone()),
-            Self::DeprecationNotice(_) | Self::ThreadRolledBack { .. } => None,
+            Self::ThreadRolledBack { thread_id, .. } => Some(thread_id.clone()),
+            Self::DeprecationNotice(_) => None,
         }
     }