agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(app-server): thread/rollback API (#8454)

Browse Source 
Add `thread/rollback` to app-server to support IDEs undo-ing the last N
turns of a thread.

For context, an IDE partner will be supporting an "undo" capability
where the IDE (the app-server client) will be responsible for reverting
the local changes made during the last turn. To support this well, we
also need a way to drop the last turn (or more generally, the last N
turns) from the agent's context. This is what `thread/rollback` does.

**Core idea**: A Thread rollback is represented as a persisted event
message (EventMsg::ThreadRollback) in the rollout JSONL file, not by
rewriting history. On resume, both the model's context (core replay) and
the UI turn list (app-server v2's thread history builder) apply these
markers so the pruned history is consistent across live conversations
and `thread/resume`.

Implementation notes:
- Rollback only affects agent context and appends to the rollout file;
clients are responsible for reverting files on disk.
- If a thread rollback is currently in progress, subsequent
`thread/rollback` calls are rejected.
- Because we use `CodexConversation::submit` and codex core tracks
active turns, returning an error on concurrent rollbacks is communicated
via an `EventMsg::Error` with a new variant
`CodexErrorInfo::ThreadRollbackFailed`. app-server watches for that and
sends the BAD_REQUEST RPC response.

Tests cover thread rollbacks in both core and app-server, including when
`num_turns` > existing turns (which clears all turns).

**Note**: this explicitly does **not** behave like `/undo` which we just
removed from the CLI, which does the opposite of what `thread/rollback`
does. `/undo` reverts local changes via ghost commits/snapshots and does
not modify the agent's context / conversation history.
This commit is contained in:
Owen Lin
2026-01-06 21:23:48 +00:00
committed by GitHub
parent 188f79afee
commit 8b7ec31ba7
21 changed files with 1187 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/context_manager/history.rs
+86
View File
@@ -5,6 +5,9 @@ use crate::truncate::approx_token_count;
use crate::truncate::approx_tokens_from_byte_count;
use crate::truncate::truncate_function_output_items_with_policy;
use crate::truncate::truncate_text;
use crate::user_instructions::SkillInstructions;
use crate::user_instructions::UserInstructions;
use crate::user_shell_command::is_user_shell_command_text;
use codex_protocol::models::ContentItem;
use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::FunctionCallOutputPayload;
@@ -152,6 +155,39 @@ impl ContextManager {
}
}
/// Drop the last `num_turns` user turns from this history.
///
/// "User turns" are identified as `ResponseItem::Message` entries whose role is `"user"`.
///
/// This mirrors thread-rollback semantics:
/// - `num_turns == 0` is a no-op
/// - if there are no user turns, this is a no-op
/// - if `num_turns` exceeds the number of user turns, all user turns are dropped while
/// preserving any items that occurred before the first user message.
pub(crate) fn drop_last_n_user_turns(&mut self, num_turns: u32) {
if num_turns == 0 {
return;
}
// Keep behavior consistent with call sites that previously operated on `get_history()`:
// normalize first (call/output invariants), then truncate based on the normalized view.
let snapshot = self.get_history();
let user_positions = user_message_positions(&snapshot);
let Some(&first_user_idx) = user_positions.first() else {
self.replace(snapshot);
return;
};
let n_from_end = usize::try_from(num_turns).unwrap_or(usize::MAX);
let cut_idx = if n_from_end >= user_positions.len() {
first_user_idx
} else {
user_positions[user_positions.len() - n_from_end]
};
self.replace(snapshot[..cut_idx].to_vec());
}
pub(crate) fn update_token_info(
&mut self,
usage: &TokenUsage,
@@ -291,6 +327,56 @@ fn estimate_reasoning_length(encoded_len: usize) -> usize {
.saturating_sub(650)
}
fn is_session_prefix(text: &str) -> bool {
let trimmed = text.trim_start();
let lowered = trimmed.to_ascii_lowercase();
lowered.starts_with("<environment_context>")
}
fn is_user_turn_boundary(item: &ResponseItem) -> bool {
let ResponseItem::Message { role, content, .. } = item else {
return false;
};
if role != "user" {
return false;
}
if UserInstructions::is_user_instructions(content)
|| SkillInstructions::is_skill_instructions(content)
{
return false;
}
for content_item in content {
match content_item {
ContentItem::InputText { text } => {
if is_session_prefix(text) || is_user_shell_command_text(text) {
return false;
}
}
ContentItem::OutputText { text } => {
if is_session_prefix(text) {
return false;
}
}
ContentItem::InputImage { .. } => {}
}
}
true
}
fn user_message_positions(items: &[ResponseItem]) -> Vec<usize> {
let mut positions = Vec::new();
for (idx, item) in items.iter().enumerate() {
if is_user_turn_boundary(item) {
positions.push(idx);
}
}
positions
}
#[cfg(test)]
#[path = "history_tests.rs"]
mod tests;
codex-rs/core/src/context_manager/history_tests.rs
+131
View File
@@ -43,6 +43,16 @@ fn user_msg(text: &str) -> ResponseItem {
}
}
fn user_input_text_msg(text: &str) -> ResponseItem {
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: text.to_string(),
}],
}
}
fn reasoning_msg(text: &str) -> ResponseItem {
ResponseItem::Reasoning {
id: String::new(),
@@ -227,6 +237,127 @@ fn remove_first_item_handles_local_shell_pair() {
assert_eq!(h.contents(), vec![]);
}
#[test]
fn drop_last_n_user_turns_preserves_prefix() {
let items = vec![
assistant_msg("session prefix item"),
user_msg("u1"),
assistant_msg("a1"),
user_msg("u2"),
assistant_msg("a2"),
];
let mut history = create_history_with_items(items);
history.drop_last_n_user_turns(1);
assert_eq!(
history.get_history(),
vec![
assistant_msg("session prefix item"),
user_msg("u1"),
assistant_msg("a1"),
]
);
let mut history = create_history_with_items(vec![
assistant_msg("session prefix item"),
user_msg("u1"),
assistant_msg("a1"),
user_msg("u2"),
assistant_msg("a2"),
]);
history.drop_last_n_user_turns(99);
assert_eq!(
history.get_history(),
vec![assistant_msg("session prefix item")]
);
}
#[test]
fn drop_last_n_user_turns_ignores_session_prefix_user_messages() {
let items = vec![
user_input_text_msg("<environment_context>ctx</environment_context>"),
user_input_text_msg("<user_instructions>do the thing</user_instructions>"),
user_input_text_msg(
"# AGENTS.md instructions for test_directory\n\n<INSTRUCTIONS>\ntest_text\n</INSTRUCTIONS>",
),
user_input_text_msg(
"<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>",
),
user_input_text_msg("<user_shell_command>echo 42</user_shell_command>"),
user_input_text_msg("turn 1 user"),
assistant_msg("turn 1 assistant"),
user_input_text_msg("turn 2 user"),
assistant_msg("turn 2 assistant"),
];
let mut history = create_history_with_items(items);
history.drop_last_n_user_turns(1);
let expected_prefix_and_first_turn = vec![
user_input_text_msg("<environment_context>ctx</environment_context>"),
user_input_text_msg("<user_instructions>do the thing</user_instructions>"),
user_input_text_msg(
"# AGENTS.md instructions for test_directory\n\n<INSTRUCTIONS>\ntest_text\n</INSTRUCTIONS>",
),
user_input_text_msg(
"<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>",
),
user_input_text_msg("<user_shell_command>echo 42</user_shell_command>"),
user_input_text_msg("turn 1 user"),
assistant_msg("turn 1 assistant"),
];
assert_eq!(history.get_history(), expected_prefix_and_first_turn);
let expected_prefix_only = vec![
user_input_text_msg("<environment_context>ctx</environment_context>"),
user_input_text_msg("<user_instructions>do the thing</user_instructions>"),
user_input_text_msg(
"# AGENTS.md instructions for test_directory\n\n<INSTRUCTIONS>\ntest_text\n</INSTRUCTIONS>",
),
user_input_text_msg(
"<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>",
),
user_input_text_msg("<user_shell_command>echo 42</user_shell_command>"),
];
let mut history = create_history_with_items(vec![
user_input_text_msg("<environment_context>ctx</environment_context>"),
user_input_text_msg("<user_instructions>do the thing</user_instructions>"),
user_input_text_msg(
"# AGENTS.md instructions for test_directory\n\n<INSTRUCTIONS>\ntest_text\n</INSTRUCTIONS>",
),
user_input_text_msg(
"<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>",
),
user_input_text_msg("<user_shell_command>echo 42</user_shell_command>"),
user_input_text_msg("turn 1 user"),
assistant_msg("turn 1 assistant"),
user_input_text_msg("turn 2 user"),
assistant_msg("turn 2 assistant"),
]);
history.drop_last_n_user_turns(2);
assert_eq!(history.get_history(), expected_prefix_only);
let mut history = create_history_with_items(vec![
user_input_text_msg("<environment_context>ctx</environment_context>"),
user_input_text_msg("<user_instructions>do the thing</user_instructions>"),
user_input_text_msg(
"# AGENTS.md instructions for test_directory\n\n<INSTRUCTIONS>\ntest_text\n</INSTRUCTIONS>",
),
user_input_text_msg(
"<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>",
),
user_input_text_msg("<user_shell_command>echo 42</user_shell_command>"),
user_input_text_msg("turn 1 user"),
assistant_msg("turn 1 assistant"),
user_input_text_msg("turn 2 user"),
assistant_msg("turn 2 assistant"),
]);
history.drop_last_n_user_turns(3);
assert_eq!(history.get_history(), expected_prefix_only);
}
#[test]
fn remove_first_item_handles_custom_tool_pair() {
let items = vec![