Skip to content

Commit 129912c

Browse files
jif-oaijif-oai
authored
feat: undo wiring (#5630)
1 parent c396a10 commit 129912c

File tree

14 files changed

+450
-65
lines changed

14 files changed

+450
-65
lines changed

‎codex-rs/Cargo.lock‎

Lines changed: 10 additions & 6 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎codex-rs/core/src/codex.rs‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ use crate::tasks::RegularTask;
109109
use crate::tasks::ReviewTask;
110110
use crate::tasks::SessionTask;
111111
use crate::tasks::SessionTaskContext;
112+
use crate::tasks::UndoTask;
112113
use crate::tools::ToolRouter;
113114
use crate::tools::context::SharedTurnDiffTracker;
114115
use crate::tools::parallel::ToolCallRuntime;
@@ -906,7 +907,7 @@ impl Session {
906907
state.record_items(items.iter());
907908
}
908909

909-
async fn replace_history(&self, items: Vec<ResponseItem>) {
910+
pub(crate) async fn replace_history(&self, items: Vec<ResponseItem>) {
910911
let mut state = self.state.lock().await;
911912
state.replace_history(items);
912913
}
@@ -1358,6 +1359,13 @@ async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiv
13581359
};
13591360
sess.send_event_raw(event).await;
13601361
}
1362+
Op::Undo => {
1363+
let turn_context = sess
1364+
.new_turn_with_sub_id(sub.id.clone(), SessionSettingsUpdate::default())
1365+
.await;
1366+
sess.spawn_task(turn_context, Vec::new(), UndoTask::new())
1367+
.await;
1368+
}
13611369
Op::Compact => {
13621370
let turn_context = sess
13631371
.new_turn_with_sub_id(sub.id.clone(), SessionSettingsUpdate::default())

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ use codex_protocol::models::ResponseItem;
1010
use codex_protocol::user_input::UserInput;
1111
use codex_utils_readiness::Readiness;
1212
use codex_utils_readiness::Token;
13-
use std::borrow::ToOwned;
1413
use std::sync::Arc;
1514
use tokio_util::sync::CancellationToken;
1615
use tracing::info;
@@ -52,8 +51,7 @@ impl SessionTask for GhostSnapshotTask {
5251
session
5352
.session
5453
.record_conversation_items(&[ResponseItem::GhostSnapshot {
55-
commit_id: ghost_commit.id().to_string(),
56-
parent: ghost_commit.parent().map(ToOwned::to_owned),
54+
ghost_commit: ghost_commit.clone(),
5755
}])
5856
.await;
5957
info!("ghost commit captured: {}", ghost_commit.id());

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ mod compact;
22
mod ghost_snapshot;
33
mod regular;
44
mod review;
5+
mod undo;
56

67
use std::sync::Arc;
78
use std::time::Duration;
@@ -29,6 +30,7 @@ pub(crate) use compact::CompactTask;
2930
pub(crate) use ghost_snapshot::GhostSnapshotTask;
3031
pub(crate) use regular::RegularTask;
3132
pub(crate) use review::ReviewTask;
33+
pub(crate) use undo::UndoTask;
3234

3335
const GRACEFULL_INTERRUPTION_TIMEOUT_MS: u64 = 100;
3436

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

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
use std::sync::Arc;
2+
3+
use crate::codex::TurnContext;
4+
use crate::protocol::EventMsg;
5+
use crate::protocol::UndoCompletedEvent;
6+
use crate::protocol::UndoStartedEvent;
7+
use crate::state::TaskKind;
8+
use crate::tasks::SessionTask;
9+
use crate::tasks::SessionTaskContext;
10+
use async_trait::async_trait;
11+
use codex_git_tooling::restore_ghost_commit;
12+
use codex_protocol::models::ResponseItem;
13+
use codex_protocol::user_input::UserInput;
14+
use tokio_util::sync::CancellationToken;
15+
use tracing::error;
16+
use tracing::info;
17+
use tracing::warn;
18+
19+
pub(crate) struct UndoTask;
20+
21+
impl UndoTask {
22+
pub(crate) fn new() -> Self {
23+
Self
24+
}
25+
}
26+
27+
#[async_trait]
28+
impl SessionTask for UndoTask {
29+
fn kind(&self) -> TaskKind {
30+
TaskKind::Regular
31+
}
32+
33+
async fn run(
34+
self: Arc<Self>,
35+
session: Arc<SessionTaskContext>,
36+
ctx: Arc<TurnContext>,
37+
_input: Vec<UserInput>,
38+
cancellation_token: CancellationToken,
39+
) -> Option<String> {
40+
let sess = session.clone_session();
41+
sess.send_event(
42+
ctx.as_ref(),
43+
EventMsg::UndoStarted(UndoStartedEvent {
44+
message: Some("Undo in progress...".to_string()),
45+
}),
46+
)
47+
.await;
48+
49+
if cancellation_token.is_cancelled() {
50+
sess.send_event(
51+
ctx.as_ref(),
52+
EventMsg::UndoCompleted(UndoCompletedEvent {
53+
success: false,
54+
message: Some("Undo cancelled.".to_string()),
55+
}),
56+
)
57+
.await;
58+
return None;
59+
}
60+
61+
let mut history = sess.clone_history().await;
62+
let mut items = history.get_history();
63+
let mut completed = UndoCompletedEvent {
64+
success: false,
65+
message: None,
66+
};
67+
68+
let Some((idx, ghost_commit)) =
69+
items
70+
.iter()
71+
.enumerate()
72+
.rev()
73+
.find_map(|(idx, item)| match item {
74+
ResponseItem::GhostSnapshot { ghost_commit } => {
75+
Some((idx, ghost_commit.clone()))
76+
}
77+
_ => None,
78+
})
79+
else {
80+
completed.message = Some("No ghost snapshot available to undo.".to_string());
81+
sess.send_event(ctx.as_ref(), EventMsg::UndoCompleted(completed))
82+
.await;
83+
return None;
84+
};
85+
86+
let commit_id = ghost_commit.id().to_string();
87+
let repo_path = ctx.cwd.clone();
88+
let restore_result =
89+
tokio::task::spawn_blocking(move || restore_ghost_commit(&repo_path, &ghost_commit))
90+
.await;
91+
92+
match restore_result {
93+
Ok(Ok(())) => {
94+
items.remove(idx);
95+
sess.replace_history(items).await;
96+
let short_id: String = commit_id.chars().take(7).collect();
97+
info!(commit_id = commit_id, "Undo restored ghost snapshot");
98+
completed.success = true;
99+
completed.message = Some(format!("Undo restored snapshot {short_id}."));
100+
}
101+
Ok(Err(err)) => {
102+
let message = format!("Failed to restore snapshot {commit_id}: {err}");
103+
warn!("{message}");
104+
completed.message = Some(message);
105+
}
106+
Err(err) => {
107+
let message = format!("Failed to restore snapshot {commit_id}: {err}");
108+
error!("{message}");
109+
completed.message = Some(message);
110+
}
111+
}
112+
113+
sess.send_event(ctx.as_ref(), EventMsg::UndoCompleted(completed))
114+
.await;
115+
None
116+
}
117+
}

‎codex-rs/exec/src/event_processor_with_human_output.rs‎

Lines changed: 18 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ use codex_core::protocol::StreamErrorEvent;
2020
use codex_core::protocol::TaskCompleteEvent;
2121
use codex_core::protocol::TurnAbortReason;
2222
use codex_core::protocol::TurnDiffEvent;
23-
use codex_core::protocol::WebSearchBeginEvent;
2423
use codex_core::protocol::WebSearchEndEvent;
2524
use codex_protocol::num_format::format_with_separators;
2625
use owo_colors::OwoColorize;
@@ -216,7 +215,6 @@ impl EventProcessor for EventProcessorWithHumanOutput {
216215
cwd.to_string_lossy(),
217216
);
218217
}
219-
EventMsg::ExecCommandOutputDelta(_) => {}
220218
EventMsg::ExecCommandEnd(ExecCommandEndEvent {
221219
aggregated_output,
222220
duration,
@@ -283,7 +281,6 @@ impl EventProcessor for EventProcessorWithHumanOutput {
283281
}
284282
}
285283
}
286-
EventMsg::WebSearchBegin(WebSearchBeginEvent { call_id: _ }) => {}
287284
EventMsg::WebSearchEnd(WebSearchEndEvent { call_id: _, query }) => {
288285
ts_msg!(self, "🌐 Searched: {query}");
289286
}
@@ -411,12 +408,6 @@ impl EventProcessor for EventProcessorWithHumanOutput {
411408
);
412409
eprintln!("{unified_diff}");
413410
}
414-
EventMsg::ExecApprovalRequest(_) => {
415-
// Should we exit?
416-
}
417-
EventMsg::ApplyPatchApprovalRequest(_) => {
418-
// Should we exit?
419-
}
420411
EventMsg::AgentReasoning(agent_reasoning_event) => {
421412
if self.show_agent_reasoning {
422413
ts_msg!(
@@ -481,15 +472,6 @@ impl EventProcessor for EventProcessorWithHumanOutput {
481472
}
482473
}
483474
}
484-
EventMsg::GetHistoryEntryResponse(_) => {
485-
// Currently ignored in exec output.
486-
}
487-
EventMsg::McpListToolsResponse(_) => {
488-
// Currently ignored in exec output.
489-
}
490-
EventMsg::ListCustomPromptsResponse(_) => {
491-
// Currently ignored in exec output.
492-
}
493475
EventMsg::ViewImageToolCall(view) => {
494476
ts_msg!(
495477
self,
@@ -510,15 +492,24 @@ impl EventProcessor for EventProcessorWithHumanOutput {
510492
}
511493
},
512494
EventMsg::ShutdownComplete => return CodexStatus::Shutdown,
513-
EventMsg::ConversationPath(_) => {}
514-
EventMsg::UserMessage(_) => {}
515-
EventMsg::EnteredReviewMode(_) => {}
516-
EventMsg::ExitedReviewMode(_) => {}
517-
EventMsg::AgentMessageDelta(_) => {}
518-
EventMsg::AgentReasoningDelta(_) => {}
519-
EventMsg::AgentReasoningRawContentDelta(_) => {}
520-
EventMsg::ItemStarted(_) => {}
521-
EventMsg::ItemCompleted(_) => {}
495+
EventMsg::WebSearchBegin(_)
496+
| EventMsg::ExecApprovalRequest(_)
497+
| EventMsg::ApplyPatchApprovalRequest(_)
498+
| EventMsg::ExecCommandOutputDelta(_)
499+
| EventMsg::GetHistoryEntryResponse(_)
500+
| EventMsg::McpListToolsResponse(_)
501+
| EventMsg::ListCustomPromptsResponse(_)
502+
| EventMsg::ConversationPath(_)
503+
| EventMsg::UserMessage(_)
504+
| EventMsg::EnteredReviewMode(_)
505+
| EventMsg::ExitedReviewMode(_)
506+
| EventMsg::AgentMessageDelta(_)
507+
| EventMsg::AgentReasoningDelta(_)
508+
| EventMsg::AgentReasoningRawContentDelta(_)
509+
| EventMsg::ItemStarted(_)
510+
| EventMsg::ItemCompleted(_)
511+
| EventMsg::UndoCompleted(_)
512+
| EventMsg::UndoStarted(_) => {}
522513
}
523514
CodexStatus::Running
524515
}

‎codex-rs/git-tooling/Cargo.toml‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,20 @@ name = "codex_git_tooling"
99
path = "src/lib.rs"
1010

1111
[dependencies]
12-
tempfile = "3"
13-
thiserror = "2"
14-
walkdir = "2"
12+
tempfile = { workspace = true }
13+
thiserror = { workspace = true }
14+
walkdir = { workspace = true }
15+
schemars = { workspace = true }
16+
serde = { workspace = true, features = ["derive"] }
17+
ts-rs = { workspace = true, features = [
18+
"uuid-impl",
19+
"serde-json-impl",
20+
"no-serde-warnings",
21+
] }
1522

1623
[lints]
1724
workspace = true
1825

1926
[dev-dependencies]
2027
assert_matches = { workspace = true }
21-
pretty_assertions = "1.4.1"
28+
pretty_assertions = { workspace = true }

0 commit comments

Comments
 (0)