diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 9bc52576d5..9654557b77 100644 --- a/codex-rs/tui/src/app.rs +++ b/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>); diff --git a/codex-rs/tui/src/app/app_server_adapter.rs b/codex-rs/tui/src/app/app_server_adapter.rs index 9abf4bf9ac..474aa5e2c3 100644 --- a/codex-rs/tui/src/app/app_server_adapter.rs +++ b/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; diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 74b050197f..2616bdd2f7 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/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 diff --git a/codex-rs/tui/src/thread_update.rs b/codex-rs/tui/src/thread_update.rs index 64b8a7838a..509b7f199c 100644 --- a/codex-rs/tui/src/thread_update.rs +++ b/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, } }