feat: TUI undo op (#5629)

codex-rs/core/src/codex.rs

@@ -109,6 +109,7 @@ use crate::tasks::RegularTask;
 use crate::tasks::ReviewTask;
 use crate::tasks::SessionTask;
 use crate::tasks::SessionTaskContext;
+use crate::tasks::UndoTask;
 use crate::tools::ToolRouter;
 use crate::tools::context::SharedTurnDiffTracker;
 use crate::tools::parallel::ToolCallRuntime;
@@ -958,7 +959,7 @@ impl Session {
         state.record_items(items.iter());
     }
 
-    async fn replace_history(&self, items: Vec<ResponseItem>) {
+    pub(crate) async fn replace_history(&self, items: Vec<ResponseItem>) {
         let mut state = self.state.lock().await;
         state.replace_history(items);
     }
@@ -1420,6 +1421,13 @@ async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiv
                 };
                 sess.send_event_raw(event).await;
             }
+            Op::Undo => {
+                let turn_context = sess
+                    .new_turn_with_sub_id(sub.id.clone(), SessionSettingsUpdate::default())
+                    .await;
+                sess.spawn_task(turn_context, Vec::new(), UndoTask::new())
+                    .await;
+            }
             Op::Compact => {
                 let turn_context = sess
                     .new_turn_with_sub_id(sub.id.clone(), SessionSettingsUpdate::default())

codex-rs/core/src/rollout/policy.rs

@@ -43,6 +43,7 @@ pub(crate) fn should_persist_event_msg(ev: &EventMsg) -> bool {
         | EventMsg::TokenCount(_)
         | EventMsg::EnteredReviewMode(_)
         | EventMsg::ExitedReviewMode(_)
+        | EventMsg::UndoCompleted(_)
         | EventMsg::TurnAborted(_) => true,
         EventMsg::Error(_)
         | EventMsg::TaskStarted(_)
@@ -68,6 +69,7 @@ pub(crate) fn should_persist_event_msg(ev: &EventMsg) -> bool {
         | EventMsg::PatchApplyEnd(_)
         | EventMsg::TurnDiff(_)
         | EventMsg::GetHistoryEntryResponse(_)
+        | EventMsg::UndoStarted(_)
         | EventMsg::McpListToolsResponse(_)
         | EventMsg::ListCustomPromptsResponse(_)
         | EventMsg::PlanUpdate(_)

codex-rs/core/src/tasks/ghost_snapshot.rs

@@ -10,7 +10,6 @@ use codex_protocol::models::ResponseItem;
 use codex_protocol::user_input::UserInput;
 use codex_utils_readiness::Readiness;
 use codex_utils_readiness::Token;
-use std::borrow::ToOwned;
 use std::sync::Arc;
 use tokio_util::sync::CancellationToken;
 use tracing::info;
@@ -52,8 +51,7 @@ impl SessionTask for GhostSnapshotTask {
                             session
                                 .session
                                 .record_conversation_items(&ctx, &[ResponseItem::GhostSnapshot {
-                                    commit_id: ghost_commit.id().to_string(),
-                                    parent: ghost_commit.parent().map(ToOwned::to_owned),
+                                    ghost_commit: ghost_commit.clone(),
                                 }])
                                 .await;
                             info!("ghost commit captured: {}", ghost_commit.id());

codex-rs/core/src/tasks/mod.rs

@@ -2,6 +2,7 @@ mod compact;
 mod ghost_snapshot;
 mod regular;
 mod review;
+mod undo;
 
 use std::sync::Arc;
 use std::time::Duration;
@@ -29,6 +30,7 @@ pub(crate) use compact::CompactTask;
 pub(crate) use ghost_snapshot::GhostSnapshotTask;
 pub(crate) use regular::RegularTask;
 pub(crate) use review::ReviewTask;
+pub(crate) use undo::UndoTask;
 
 const GRACEFULL_INTERRUPTION_TIMEOUT_MS: u64 = 100;
 

codex-rs/core/src/tasks/undo.rs

@@ -0,0 +1,117 @@
+use std::sync::Arc;
+
+use crate::codex::TurnContext;
+use crate::protocol::EventMsg;
+use crate::protocol::UndoCompletedEvent;
+use crate::protocol::UndoStartedEvent;
+use crate::state::TaskKind;
+use crate::tasks::SessionTask;
+use crate::tasks::SessionTaskContext;
+use async_trait::async_trait;
+use codex_git_tooling::restore_ghost_commit;
+use codex_protocol::models::ResponseItem;
+use codex_protocol::user_input::UserInput;
+use tokio_util::sync::CancellationToken;
+use tracing::error;
+use tracing::info;
+use tracing::warn;
+
+pub(crate) struct UndoTask;
+
+impl UndoTask {
+    pub(crate) fn new() -> Self {
+        Self
+    }
+}
+
+#[async_trait]
+impl SessionTask for UndoTask {
+    fn kind(&self) -> TaskKind {
+        TaskKind::Regular
+    }
+
+    async fn run(
+        self: Arc<Self>,
+        session: Arc<SessionTaskContext>,
+        ctx: Arc<TurnContext>,
+        _input: Vec<UserInput>,
+        cancellation_token: CancellationToken,
+    ) -> Option<String> {
+        let sess = session.clone_session();
+        sess.send_event(
+            ctx.as_ref(),
+            EventMsg::UndoStarted(UndoStartedEvent {
+                message: Some("Undo in progress...".to_string()),
+            }),
+        )
+        .await;
+
+        if cancellation_token.is_cancelled() {
+            sess.send_event(
+                ctx.as_ref(),
+                EventMsg::UndoCompleted(UndoCompletedEvent {
+                    success: false,
+                    message: Some("Undo cancelled.".to_string()),
+                }),
+            )
+            .await;
+            return None;
+        }
+
+        let mut history = sess.clone_history().await;
+        let mut items = history.get_history();
+        let mut completed = UndoCompletedEvent {
+            success: false,
+            message: None,
+        };
+
+        let Some((idx, ghost_commit)) =
+            items
+                .iter()
+                .enumerate()
+                .rev()
+                .find_map(|(idx, item)| match item {
+                    ResponseItem::GhostSnapshot { ghost_commit } => {
+                        Some((idx, ghost_commit.clone()))
+                    }
+                    _ => None,
+                })
+        else {
+            completed.message = Some("No ghost snapshot available to undo.".to_string());
+            sess.send_event(ctx.as_ref(), EventMsg::UndoCompleted(completed))
+                .await;
+            return None;
+        };
+
+        let commit_id = ghost_commit.id().to_string();
+        let repo_path = ctx.cwd.clone();
+        let restore_result =
+            tokio::task::spawn_blocking(move || restore_ghost_commit(&repo_path, &ghost_commit))
+                .await;
+
+        match restore_result {
+            Ok(Ok(())) => {
+                items.remove(idx);
+                sess.replace_history(items).await;
+                let short_id: String = commit_id.chars().take(7).collect();
+                info!(commit_id = commit_id, "Undo restored ghost snapshot");
+                completed.success = true;
+                completed.message = Some(format!("Undo restored snapshot {short_id}."));
+            }
+            Ok(Err(err)) => {
+                let message = format!("Failed to restore snapshot {commit_id}: {err}");
+                warn!("{message}");
+                completed.message = Some(message);
+            }
+            Err(err) => {
+                let message = format!("Failed to restore snapshot {commit_id}: {err}");
+                error!("{message}");
+                completed.message = Some(message);
+            }
+        }
+
+        sess.send_event(ctx.as_ref(), EventMsg::UndoCompleted(completed))
+            .await;
+        None
+    }
+}