openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 17:38:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Preserve acceptance order in retained thread context (#42770)

Browse Source 
## Why

Queued prompts and `request_user_input` answers can be persisted in a different
order from when the host accepted them. Using persistence order can therefore
reorder retained instructions during replay or apply the wrong rollback
boundary.

## What changed

- Reserve an acceptance sequence for user prompts and elicitation responses
  when retained thread context is enabled.
- Persist that sequence with user-message metadata and verified answers, and
  use it when ordering, bounding, replaying, and rolling back retained context.
- Keep events without acceptance metadata compatible with legacy recording
  order.

## Testing

- Cover delayed queued-input recording, rejected prompts, checkpoint replay,
  compaction, rollback migration, and legacy retained-context behavior.

GitOrigin-RevId: 42d24465cc2033bc51c0c2da79d3c0884a8ad98a
This commit is contained in:
felixxia-oai
2026-09-04 13:50:27 +00:00
committed by copyberry
parent 88f87d907a
commit 99d66aa1c5
24 changed files with 489 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst
Generated
BIN
View File
Binary file not shown.
codex-rs/core/src/context_manager/history.rs
+14 -1
View File
@@ -187,6 +187,10 @@ impl ContextManager {
self.retain_user_messages = true;
}
pub(crate) fn reserve_input_order(&mut self) -> u64 {
Arc::make_mut(&mut self.retained_context).reserve_order()
}
pub(crate) fn record_retained_context(&mut self, event: &RetainedContextEvent) -> bool {
if !Arc::make_mut(&mut self.retained_context).record(event) {
return false;
@@ -359,6 +363,7 @@ impl ContextManager {
text,
complete,
},
metadata.and_then(|metadata| metadata.user_input_order),
);
} else {
Arc::make_mut(&mut self.retained_context).mark_user_messages_incomplete();
@@ -536,6 +541,10 @@ impl ContextManager {
let first_removed_message_id = snapshot[cut_idx]
.id()
.map(codex_protocol::ResponseItemId::as_str);
let acceptance_order = snapshot[cut_idx]
.metadata
.as_ref()
.and_then(|metadata| metadata.user_input_order);
let mut review_history = self.review_history.take();
if let Some(history) = &mut review_history {
history.truncate_before(&snapshot[cut_idx].item);
@@ -577,7 +586,11 @@ impl ContextManager {
.filter_map(|item| item.turn_id())
.collect::<Vec<_>>();
if self.retain_user_messages {
Arc::make_mut(&mut retained_context).rollback(&removed_turns, first_removed_message_id);
Arc::make_mut(&mut retained_context).rollback(
&removed_turns,
first_removed_message_id,
acceptance_order,
);
} else {
Arc::make_mut(&mut retained_context).retain_answers(|answer| {
// Legacy answers follow their original call, not later steers in the same turn.
codex-rs/core/src/context_manager/history_tests.rs
+8 -6
View File
@@ -1252,8 +1252,8 @@ fn drop_last_n_user_turns_preserves_prefix() {
);
// With no remaining instruction boundary, rollback must not revoke facts from a
// prior checkpoint merely because their source messages are no longer visible.
history.record_retained_context(&codex_history::RetainedContextEvent::VerifiedAnswer(
codex_history::VerifiedAnswer {
history.record_retained_context(&codex_history::RetainedContextEvent::VerifiedAnswer {
answer: codex_history::VerifiedAnswer {
turn_id: "checkpoint-turn".to_owned(),
call_id: "ask-1".to_owned(),
questions: vec![codex_history::VerifiedQuestionAnswer {
@@ -1261,7 +1261,8 @@ fn drop_last_n_user_turns_preserves_prefix() {
answer: "Only privately.".to_owned(),
}],
},
));
acceptance_order: None,
});
let retained = history.retained_context().clone();
history.drop_last_n_user_turns(/*num_turns*/ 1);
assert_eq!(history.retained_context(), &retained);
@@ -1292,8 +1293,8 @@ fn drop_last_n_user_turns_preserves_prefix() {
};
history.record_items([&message], TruncationPolicy::Tokens(10_000));
if id == "restriction" {
history.record_retained_context(&codex_history::RetainedContextEvent::VerifiedAnswer(
codex_history::VerifiedAnswer {
history.record_retained_context(&codex_history::RetainedContextEvent::VerifiedAnswer {
answer: codex_history::VerifiedAnswer {
turn_id: "shared-turn".to_owned(),
call_id: "ask".to_owned(),
questions: vec![codex_history::VerifiedQuestionAnswer {
@@ -1301,7 +1302,8 @@ fn drop_last_n_user_turns_preserves_prefix() {
answer: "Only privately.".to_owned(),
}],
},
));
acceptance_order: None,
});
expected = Some(history.retained_context().clone());
}
}
codex-rs/core/src/hook_runtime.rs
+6 -1
View File
@@ -705,11 +705,16 @@ pub(crate) async fn record_pending_input(
persist_context: PersistContext,
) {
match pending_input {
TurnInput::UserInput { content, client_id } => {
TurnInput::UserInput {
content,
client_id,
acceptance_order,
} => {
sess.record_user_prompt_and_emit_turn_item(
turn_context.as_ref(),
content.as_slice(),
client_id,
acceptance_order,
persist_context,
)
.await;
codex-rs/core/src/mcp_skill_dependencies.rs
+1 -1
View File
@@ -311,7 +311,7 @@ async fn should_install_mcp_dependencies(
sess.notify_user_input_response(sub_id, empty.clone()).await;
empty
}
response = response_fut => response.unwrap_or_else(|| RequestUserInputResponse {
response = response_fut => response.map(|accepted| accepted.response).unwrap_or_else(|| RequestUserInputResponse {
answers: HashMap::new(),
}),
};
codex-rs/core/src/mcp_tool_call.rs
+1 -1
View File
@@ -1620,7 +1620,7 @@ pub(crate) async fn request_mcp_tool_user_approval(
.request_user_input(turn_context, call_id.to_string(), args)
.await;
normalize_approval_decision_for_mode(
parse_mcp_tool_approval_response(response, &question_id),
parse_mcp_tool_approval_response(response.map(|accepted| accepted.response), &question_id),
*approval_mode,
)
}
codex-rs/core/src/session/input_queue.rs
+4
View File
@@ -23,6 +23,8 @@ pub enum TurnInput {
UserInput {
content: Vec<UserInput>,
client_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
acceptance_order: Option<u64>,
},
FunctionCallOutput(ResponseItem),
// Preserve the existing serialized format while carrying injection API metadata
@@ -487,6 +489,7 @@ mod tests {
.extend_pending_input_and_accept_mailbox_delivery_for_turn_state(
&turn_state,
vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "steer".to_string(),
text_elements: Vec::new(),
@@ -519,6 +522,7 @@ mod tests {
.extend_pending_input_and_accept_mailbox_delivery_for_turn_state(
&turn_state,
vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "already pending".to_string(),
text_elements: Vec::new(),
codex-rs/core/src/session/mod.rs
+28 -8
View File
@@ -293,6 +293,7 @@ use crate::network_policy_decision::execpolicy_network_rule_amendment;
use crate::rollout::map_session_init_error;
use crate::session_startup_prewarm::SessionStartupPrewarmHandle;
use crate::shell;
use crate::state::AcceptedUserInputResponse;
use crate::state::AutoCompactWindowIds;
use crate::state::AutoCompactWindowSnapshot;
use crate::state::PendingRequestPermissions;
@@ -319,6 +320,7 @@ use codex_core_plugins::PluginCommandAttribution;
use codex_core_plugins::PluginsManager;
use codex_core_plugins::RecommendedPluginCandidatesInput;
use codex_git_utils::get_git_repo_root;
use codex_history::CodexHarnessMetadata;
use codex_history::CompactedItem;
use codex_history::InitialHistory;
use codex_history::ResponseItemEnvelope;
@@ -2947,12 +2949,12 @@ impl Session {
clippy::await_holding_invalid_type,
reason = "active turn checks and turn state updates must remain atomic"
)]
pub async fn request_user_input(
pub(crate) async fn request_user_input(
&self,
turn_context: &TurnContext,
call_id: String,
args: RequestUserInputArgs,
) -> Option<RequestUserInputResponse> {
) -> Option<AcceptedUserInputResponse> {
let _elicitation = self.services.elicitations.register();
let sub_id = turn_context.sub_id.clone();
let (tx_response, rx_response) = oneshot::channel();
@@ -2998,15 +3000,23 @@ impl Session {
let mut active = self.active_turn.lock().await;
match active.as_mut() {
Some(at) => {
let mut ts = at.turn_state.lock().await;
ts.remove_pending_user_input(sub_id)
let sender = at.turn_state.lock().await.remove_pending_user_input(sub_id);
match sender {
Some(sender) => Some((sender, self.reserve_user_input_order().await)),
None => None,
}
}
None => None,
}
};
match entry {
Some(tx_response) => {
tx_response.send(response).ok();
Some((tx_response, acceptance_order)) => {
tx_response
.send(AcceptedUserInputResponse {
response,
acceptance_order,
})
.ok();
}
None => {
warn!("No pending user input found for sub_id: {sub_id}");
@@ -4539,14 +4549,24 @@ impl Session {
turn_context: &TurnContext,
input: &[UserInput],
client_id: Option<String>,
acceptance_order: Option<u64>,
persist_context: PersistContext,
) {
// Persist the user message to history, but emit the turn item from `UserInput` so
// UI-only `text_elements` are preserved. `ResponseItem::Message` does not carry
// those spans, and `record_response_item_and_emit_turn_item` would drop them.
let response_item = self.response_item_from_user_input(input.to_vec());
self.record_conversation_items(turn_context, std::slice::from_ref(&response_item))
.await;
self.record_annotated_conversation_items(
turn_context,
vec![ResponseItemEnvelope {
item: response_item,
metadata: acceptance_order.map(|order| CodexHarnessMetadata {
user_input_order: Some(order),
..Default::default()
}),
}],
)
.await;
let mut user_message_item = UserMessageItem::new(input);
user_message_item.client_id = client_id;
let turn_item = TurnItem::UserMessage(user_message_item);
codex-rs/core/src/session/retained_context.rs
+10 -1
View File
@@ -1,5 +1,6 @@
//! Records host-verified facts at the same persistence boundary as context checkpoints.
//! Orders retained inputs and records host-verified facts at the context checkpoint boundary.
use codex_features::Feature;
use codex_history::RetainedContextEvent;
use codex_history::RolloutItem;
@@ -7,6 +8,14 @@ use super::Session;
use super::thread_settings;
impl Session {
/// Legacy mode neither reserves a sequence nor takes the session-state lock.
pub(crate) async fn reserve_user_input_order(&self) -> Option<u64> {
if !self.enabled(Feature::GuardianThreadContext) {
return None;
}
Some(self.state.lock().await.history.reserve_input_order())
}
pub(crate) async fn record_retained_context(&self, mut event: RetainedContextEvent) {
event.bound();
// Share the checkpoint persistence lock so a fact cannot land on the wrong side
codex-rs/core/src/session/review.rs
+1
View File
@@ -183,6 +183,7 @@ pub(super) async fn spawn_review_thread(
// Seed the child task with the review prompt as the initial user message.
let input = vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: review_prompt,
// Review prompt is synthesized; no UI element ranges to preserve.
codex-rs/core/src/session/rollout_reconstruction_tests.rs
+5 -3
View File
@@ -1478,15 +1478,17 @@ async fn reconstruct_history_preserves_legacy_compaction_count_with_session_meta
async fn reconstruct_history_legacy_compaction_without_replacement_history_does_not_inject_current_initial_context()
{
let (session, turn_context) = make_session_and_context().await;
let answer =
codex_history::RetainedContextEvent::VerifiedAnswer(codex_history::VerifiedAnswer {
let answer = codex_history::RetainedContextEvent::VerifiedAnswer {
answer: codex_history::VerifiedAnswer {
turn_id: "legacy-turn".to_owned(),
call_id: "ask-1".to_owned(),
questions: vec![codex_history::VerifiedQuestionAnswer {
question: "Upload?".to_owned(),
answer: "Only privately.".to_owned(),
}],
});
},
acceptance_order: None,
};
let mut retained = codex_history::RetainedContext::default();
retained.mark_user_messages_incomplete();
retained.record(&answer);
codex-rs/core/src/session/tests.rs
+11
View File
@@ -8154,6 +8154,7 @@ async fn spawn_task_turn_span_inherits_dispatch_trace_context() {
sess.spawn_task(
Arc::clone(&tc),
vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "hello".to_string(),
text_elements: Vec::new(),
@@ -9698,6 +9699,7 @@ async fn spawn_task_does_not_update_previous_turn_settings_for_non_run_turn_task
sess.set_previous_turn_settings(/*previous_turn_settings*/ None)
.await;
let input = vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "hello".to_string(),
text_elements: Vec::new(),
@@ -11468,6 +11470,7 @@ async fn guardian_auto_review_emits_thread_idle_after_interrupt() {
async fn guardian_helper_review_interrupts_after_three_consecutive_denials() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
let input = vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "keep turn active for helper reviews".to_string(),
text_elements: Vec::new(),
@@ -11535,6 +11538,7 @@ async fn turn_complete_flushes_terminal_event_after_delivery() {
.await;
let input = vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "complete normally".to_string(),
text_elements: Vec::new(),
@@ -11562,6 +11566,7 @@ async fn turn_aborted_flushes_terminal_event_after_delivery() {
.await;
let input = vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "interrupt me".to_string(),
text_elements: Vec::new(),
@@ -11604,6 +11609,7 @@ async fn turn_aborted_flushes_terminal_event_after_delivery() {
async fn abort_regular_task_emits_marker_before_turn_aborted() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
let input = vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "hello".to_string(),
text_elements: Vec::new(),
@@ -11645,6 +11651,7 @@ async fn abort_regular_task_emits_marker_before_turn_aborted() {
async fn abort_gracefully_emits_marker_before_turn_aborted() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
let input = vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "hello".to_string(),
text_elements: Vec::new(),
@@ -11706,6 +11713,7 @@ async fn submit_steer_only(
async fn task_finish_emits_turn_item_lifecycle_for_leftover_pending_user_input() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
let input = vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "hello".to_string(),
text_elements: Vec::new(),
@@ -12143,6 +12151,7 @@ async fn steered_input_reopens_mailbox_delivery_for_current_turn() {
(sess.input_queue.get_pending_input(&sess.active_turn).await).0,
vec![
TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "follow up".to_string(),
text_elements: Vec::new(),
@@ -12199,6 +12208,7 @@ async fn stale_defer_mailbox_delivery_does_not_override_steered_input() {
(sess.input_queue.get_pending_input(&sess.active_turn).await).0,
vec![
TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "follow up".to_string(),
text_elements: Vec::new(),
@@ -12270,6 +12280,7 @@ async fn tool_calls_reopen_mailbox_delivery_for_current_turn() {
async fn abort_review_task_emits_exited_then_aborted_and_records_history() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
let input = vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "start review".to_string(),
text_elements: Vec::new(),
codex-rs/core/src/session/turn_input.rs
+11 -8
View File
@@ -313,7 +313,7 @@ async fn start_or_steer(
}
let mut task_input = merge_additional_context_input(session, additional_context).await;
if has_explicit_input {
task_input.push(pending_turn_input(input));
task_input.push(pending_turn_input(session, input).await);
}
session
.spawn_task(turn_context, task_input, RegularTask::new())
@@ -424,7 +424,7 @@ async fn start_if_idle(
if let SubmittedTurnInput::UserInput { content, .. } = &input {
turn_context.session_telemetry.user_prompt(content);
}
task_input.push(pending_turn_input(input));
task_input.push(pending_turn_input(session, input).await);
}
TurnStartKind::Automatic => {
// Empty automatic user input resumes sampling without a new message.
@@ -433,7 +433,7 @@ async fn start_if_idle(
.input_queue
.extend_pending_input_for_turn_state(
turn_state.as_ref(),
vec![pending_turn_input(input)],
vec![pending_turn_input(session, input).await],
)
.await;
}
@@ -621,9 +621,10 @@ impl Session {
TurnInput::UserInput {
content: std::mem::take(content),
client_id: client_id.clone(),
acceptance_order: self.reserve_user_input_order().await,
}
}
input => pending_turn_input(input.clone()),
input => pending_turn_input(self, input.clone()).await,
};
pending_input.push(input);
if active_task
@@ -663,11 +664,13 @@ async fn merge_additional_context_input(
.collect()
}
fn pending_turn_input(input: SubmittedTurnInput) -> TurnInput {
async fn pending_turn_input(session: &Session, input: SubmittedTurnInput) -> TurnInput {
match input {
SubmittedTurnInput::UserInput { content, client_id } => {
TurnInput::UserInput { content, client_id }
}
SubmittedTurnInput::UserInput { content, client_id } => TurnInput::UserInput {
content,
client_id,
acceptance_order: session.reserve_user_input_order().await,
},
SubmittedTurnInput::ResponseItem(mut item)
if matches!(
&item,
codex-rs/core/src/session/turn_input_tests.rs
+2
View File
@@ -862,6 +862,7 @@ async fn steer_only_enforces_expected_turn_id() {
.spawn_task(
Arc::clone(&turn_context),
vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "hello".to_string(),
text_elements: Vec::new(),
@@ -951,6 +952,7 @@ async fn rejects_non_regular_turns() {
.spawn_task(
Arc::clone(&turn_context),
vec![TurnInput::UserInput {
acceptance_order: None,
content: vec![UserInput::Text {
text: "hello".to_string(),
text_elements: Vec::new(),
codex-rs/core/src/state/mod.rs
+1
View File
@@ -10,6 +10,7 @@ pub(crate) use auto_compact_window::AutoCompactWindowIds;
pub(crate) use auto_compact_window::AutoCompactWindowSnapshot;
pub(crate) use service::SessionServices;
pub(crate) use session::SessionState;
pub(crate) use turn::AcceptedUserInputResponse;
pub(crate) use turn::ActiveTurn;
pub(crate) use turn::MailboxDeliveryPhase;
pub(crate) use turn::PendingRequestPermissions;
codex-rs/core/src/state/turn.rs
+12 -5
View File
@@ -1,4 +1,4 @@
//! Turn-scoped state and active turn metadata scaffolding.
//! Turn-scoped state, input waiters, and active turn metadata.
use std::collections::HashMap;
use std::sync::Arc;
@@ -90,7 +90,7 @@ pub(crate) struct RunningTask {
pub(crate) struct TurnState {
pending_approvals: HashMap<String, oneshot::Sender<ReviewDecision>>,
pending_request_permissions: HashMap<String, PendingRequestPermissions>,
pending_user_input: HashMap<String, oneshot::Sender<RequestUserInputResponse>>,
pending_user_input: HashMap<String, oneshot::Sender<AcceptedUserInputResponse>>,
pending_elicitations: HashMap<(String, RequestId), oneshot::Sender<ElicitationResponse>>,
mcp_tool_approval_metadata: HashMap<String, (Option<McpInvocation>, McpToolApprovalMetadata)>,
pending_dynamic_tools: HashMap<String, oneshot::Sender<DynamicToolResponse>>,
@@ -106,6 +106,13 @@ pub(crate) struct TurnState {
pub(crate) last_known_step_context: Option<Arc<StepContext>>,
}
/// Host receipt metadata follows the response through the asynchronous tool waiter.
pub(crate) struct AcceptedUserInputResponse {
pub(crate) response: RequestUserInputResponse,
/// Absent in legacy mode, which does not reserve or persist acceptance order.
pub(crate) acceptance_order: Option<u64>,
}
pub(crate) struct PendingRequestPermissions {
pub(crate) tx_response: oneshot::Sender<RequestPermissionsResponse>,
pub(crate) requested_permissions: RequestPermissionProfile,
@@ -156,15 +163,15 @@ impl TurnState {
pub(crate) fn insert_pending_user_input(
&mut self,
key: String,
tx: oneshot::Sender<RequestUserInputResponse>,
) -> Option<oneshot::Sender<RequestUserInputResponse>> {
tx: oneshot::Sender<AcceptedUserInputResponse>,
) -> Option<oneshot::Sender<AcceptedUserInputResponse>> {
self.pending_user_input.insert(key, tx)
}
pub(crate) fn remove_pending_user_input(
&mut self,
key: &str,
) -> Option<oneshot::Sender<RequestUserInputResponse>> {
) -> Option<oneshot::Sender<AcceptedUserInputResponse>> {
self.pending_user_input.remove(key)
}
codex-rs/core/src/tools/handlers/request_user_input.rs
+11 -6
View File
@@ -85,7 +85,7 @@ impl RequestUserInputHandler {
auto_resolution_ms: None,
};
let questions = args.questions.clone();
let response = session
let accepted = session
.request_user_input(turn.as_ref(), call_id.clone(), args)
.await
.ok_or_else(|| {
@@ -94,6 +94,8 @@ impl RequestUserInputHandler {
))
})?;
let response = accepted.response;
let content = serde_json::to_string(&response).map_err(|err| {
FunctionCallError::Fatal(format!(
"failed to serialize {REQUEST_USER_INPUT_TOOL_NAME} response: {err}"
@@ -142,11 +144,14 @@ impl RequestUserInputHandler {
.collect::<Vec<_>>();
if !user_input.is_empty() {
session
.record_retained_context(RetainedContextEvent::VerifiedAnswer(VerifiedAnswer {
turn_id: turn.sub_id.clone(),
call_id,
questions: user_input,
}))
.record_retained_context(RetainedContextEvent::VerifiedAnswer {
answer: VerifiedAnswer {
turn_id: turn.sub_id.clone(),
call_id,
questions: user_input,
},
acceptance_order: accepted.acceptance_order,
})
.await;
}
}
codex-rs/core/tests/suite/guardian_retained_context.rs
+43 -25
View File
@@ -51,6 +51,7 @@ async fn record_answer(
server: &MockServer,
call_id: &str,
answer: &str,
acceptance_order: u64,
) -> Result<VerifiedAnswer> {
let question = format!("May I publish {call_id}?");
mount_sse_sequence(
@@ -109,8 +110,10 @@ async fn record_answer(
}],
};
thread.ensure_rollout_materialized().await;
let event =
RolloutItem::RetainedContext(RetainedContextEvent::VerifiedAnswer(retained.clone()));
let event = RolloutItem::RetainedContext(RetainedContextEvent::VerifiedAnswer {
answer: retained.clone(),
acceptance_order: Some(acceptance_order),
});
// Repeated delivery must not duplicate evidence during replay.
thread.append_rollout_items(&[event.clone(), event]).await?;
Ok(retained)
@@ -299,6 +302,9 @@ async fn retained_instructions_keep_identity_across_compaction_and_resume(
matches!(event, EventMsg::TurnComplete(_))
})
.await;
// Rebuild from source events before there is a compaction checkpoint. The answer
// was persisted first, but the accepted steering instruction must retain order 1.
let thread = resume(&test, &test.codex).await?;
assert_eq!(answers[0].turn_id, answers[1].turn_id);
let requests = response_mock.requests();
assert_eq!(requests.len(), 3);
@@ -316,7 +322,7 @@ async fn retained_instructions_keep_identity_across_compaction_and_resume(
);
}
}
let history = test.codex.conversation_history_snapshot().await;
let history = thread.conversation_history_snapshot().await;
let user_messages = [INITIAL, STEER]
.into_iter()
.enumerate()
@@ -336,7 +342,7 @@ async fn retained_instructions_keep_identity_across_compaction_and_resume(
.and_then(ResponseItem::id)
.expect("original user-message identity");
json!({
"order": index * 2, "turn_id": answers[index].turn_id,
"order": index, "turn_id": answers[index].turn_id,
"message_id": message_id.as_str(), "text": text, "complete": true,
})
})
@@ -346,7 +352,7 @@ async fn retained_instructions_keep_identity_across_compaction_and_resume(
.enumerate()
.map(|(index, answer)| {
let mut value = json!(answer);
value["order"] = json!(index * 2 + 1);
value["order"] = json!(index + 2);
value
})
.collect::<Vec<_>>();
@@ -360,19 +366,24 @@ async fn retained_instructions_keep_identity_across_compaction_and_resume(
"verified_answers": [], "incomplete": false, "next_order": 0,
});
answers.clear();
assert!(
load_context(&test, &test.codex)
.await?
.iter()
.all(|item| !matches!(item, RolloutItem::RetainedContext(_)))
);
for item in load_context(&test, &thread).await? {
assert!(!matches!(item, RolloutItem::RetainedContext(_)));
if let RolloutItem::ResponseItem(envelope) = item {
assert!(
envelope
.metadata
.as_ref()
.is_none_or(|metadata| metadata.user_input_order.is_none())
);
}
}
}
assert_eq!(serde_json::to_value(history.retained_context())?, expected);
assert_eq!(
serde_json::to_value(compact_and_assert_answers(&test, &test.codex, &answers).await?)?,
serde_json::to_value(compact_and_assert_answers(&test, &thread, &answers).await?)?,
expected
);
let compacted = test.codex.conversation_history_snapshot().await;
let compacted = thread.conversation_history_snapshot().await;
for message in &user_messages {
assert_eq!(
compacted
@@ -382,7 +393,7 @@ async fn retained_instructions_keep_identity_across_compaction_and_resume(
"only thread-owned context preserves original user-message identity"
);
}
let mut thread = Arc::clone(&test.codex);
let mut thread = thread;
for &remaining in rollback_counts {
let remaining = if thread_context_enabled { remaining } else { 0 };
thread = resume(&test, &thread).await?;
@@ -407,11 +418,9 @@ async fn retained_instructions_keep_identity_across_compaction_and_resume(
expected["verified_answers"]
.as_array_mut()
.expect("expected retained verified answers")
.truncate(remaining);
.clear();
assert_eq!(
serde_json::to_value(
compact_and_assert_answers(&test, &thread, &answers[..remaining]).await?
)?,
serde_json::to_value(compact_and_assert_answers(&test, &thread, &[]).await?)?,
expected
);
}
@@ -537,11 +546,12 @@ async fn retained_answers_rollback_only_the_steered_instruction() -> Result<()>
.with_history_mode(ThreadHistoryMode::Legacy)
.with_config(|config| {
config.experimental_thread_store = ThreadStoreConfig::Local;
for feature in [
Feature::TokenBudget,
Feature::DefaultModeRequestUserInput,
Feature::GuardianThreadContext,
] {
// Legacy saved answers use their source calls, not acceptance-order metadata.
config
.features
.disable(Feature::GuardianThreadContext)
.expect("legacy evidence fixture");
for feature in [Feature::TokenBudget, Feature::DefaultModeRequestUserInput] {
config
.features
.enable(feature)
@@ -653,7 +663,10 @@ async fn retained_answers_rollback_only_the_steered_instruction() -> Result<()>
.iter()
.cloned()
.map(|answer| {
RolloutItem::RetainedContext(RetainedContextEvent::VerifiedAnswer(answer))
RolloutItem::RetainedContext(RetainedContextEvent::VerifiedAnswer {
answer,
acceptance_order: None,
})
})
.collect::<Vec<_>>(),
)
@@ -714,6 +727,7 @@ async fn retained_answers_cross_real_session_boundaries(
&server,
"before-compact",
"Only publish privately.",
/*acceptance_order*/ 1,
)
.await?;
let thread = resume(&test, &test.codex).await?;
@@ -724,6 +738,7 @@ async fn retained_answers_cross_real_session_boundaries(
&server,
"after-compact",
"Do not publish after all.",
/*acceptance_order*/ 3,
)
.await?;
thread.flush_rollout().await?;
@@ -756,7 +771,10 @@ async fn retained_answers_cross_real_session_boundaries(
events.dedup();
assert_eq!(
events,
vec![RetainedContextEvent::VerifiedAnswer(after.clone())]
vec![RetainedContextEvent::VerifiedAnswer {
answer: after.clone(),
acceptance_order: Some(3)
}]
);
let thread = resume(&test, &thread).await?;
let expected = [before.clone(), after];
codex-rs/core/tests/suite/hooks.rs
+36 -2
View File
@@ -70,6 +70,7 @@ use core_test_support::wait_for_event;
use core_test_support::wait_for_mcp_server;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
use std::sync::Arc;
use std::time::Duration;
use tempfile::TempDir;
@@ -2649,8 +2650,12 @@ async fn blocked_user_prompt_submit_persists_additional_context_for_next_turn()
Ok(())
}
#[test_case::test_case(/*thread_context_enabled*/ true; "retained context enabled")]
#[test_case::test_case(/*thread_context_enabled*/ false; "retained context disabled")]
#[tokio::test]
async fn blocked_queued_prompt_does_not_strand_earlier_accepted_prompt() -> Result<()> {
async fn blocked_queued_prompt_does_not_strand_earlier_accepted_prompt(
thread_context_enabled: bool,
) -> Result<()> {
skip_if_no_network!(Ok(()));
let (gate_completed_tx, gate_completed_rx) = oneshot::channel();
@@ -2693,7 +2698,13 @@ async fn blocked_queued_prompt_does_not_strand_earlier_accepted_prompt() -> Resu
write_user_prompt_submit_hook(home, "blocked queued prompt", BLOCKED_PROMPT_CONTEXT)
.expect("failed to write user prompt submit hook test fixture");
})
.with_config(trust_discovered_hooks);
.with_config(move |config| {
trust_discovered_hooks(config);
config
.features
.set_enabled(Feature::GuardianThreadContext, thread_context_enabled)
.expect("test context mode");
});
let test = builder.build_with_streaming_server(&server).await?;
test.codex
@@ -2748,6 +2759,29 @@ async fn blocked_queued_prompt_does_not_strand_earlier_accepted_prompt() -> Resu
"second request should not include the blocked queued prompt",
);
let history = test.codex.conversation_history_snapshot().await;
let retained = serde_json::to_value(history.retained_context())?;
assert_eq!(
retained["user_messages"]
.as_array()
.expect("retained user messages")
.iter()
.map(|message| (message["order"].clone(), message["text"].clone()))
.collect::<Vec<_>>(),
if thread_context_enabled {
vec![
(json!(0), json!("initial prompt")),
(json!(1), json!("accepted queued prompt")),
]
} else {
Vec::new()
},
);
assert_eq!(
retained["next_order"],
json!(if thread_context_enabled { 3 } else { 0 })
);
let hook_inputs = read_user_prompt_submit_hook_inputs(test.codex_home_path())?;
assert_eq!(hook_inputs.len(), 3);
assert_eq!(
codex-rs/history/src/lib.rs
+4
View File
@@ -62,6 +62,10 @@ pub struct CodexHarnessMetadata {
/// Producer compatibility for an opaque compaction item, never the currently selected model.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub compaction_model_hash: Option<String>,
/// Thread acceptance order, independent of when queued user input reaches model history.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub user_input_order: Option<u64>,
}
impl ResponseItemEnvelope {
codex-rs/history/src/retained_context.rs
+60 -19
View File
@@ -73,11 +73,17 @@ impl std::fmt::Debug for VerifiedAnswer {
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum RetainedContextEvent {
VerifiedAnswer(VerifiedAnswer),
VerifiedAnswer {
#[serde(flatten)]
answer: VerifiedAnswer,
/// Absent in legacy events, which retain their recorded ordering.
#[serde(default, skip_serializing_if = "Option::is_none")]
acceptance_order: Option<u64>,
},
}
/// Bounded snapshot of retained families, persisted with the parent compaction checkpoint.
/// Facts live until their source instruction is rolled back; compaction does not expire them.
/// Facts live until their instruction boundary is rolled back; compaction does not expire them.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, JsonSchema)]
pub struct RetainedContext {
verified_answers: VecDeque<Ordered<VerifiedAnswer>>,
@@ -98,6 +104,9 @@ fn legacy_user_messages_incomplete() -> bool {
}
fn bound_family<T: Serialize>(items: &mut VecDeque<Ordered<T>>, incomplete: &mut bool) {
// Queued instructions can be recorded after later-accepted answers. Evict by
// acceptance order, not by the order in which persistence happened to finish.
items.make_contiguous().sort_by_key(|entry| entry.order);
while items.len() > MAX_FAMILY_RECORDS
|| serde_json::to_vec(items).map_or(true, |bytes| bytes.len() > MAX_FAMILY_BYTES)
{
@@ -124,7 +133,7 @@ impl RetainedContextEvent {
/// Bounds a persisted event before it enters the rollout or the live snapshot.
pub fn bound(&mut self) {
match self {
Self::VerifiedAnswer(answer) => {
Self::VerifiedAnswer { answer, .. } => {
if serde_json::to_vec(answer).map_or(true, |bytes| bytes.len() > MAX_RECORD_BYTES) {
answer.questions.clear();
// IDs are correlation metadata, not model-visible authorization text.
@@ -141,6 +150,19 @@ impl RetainedContextEvent {
}
impl RetainedContext {
/// Reserves order without retaining pending input that hooks may reject or cancel.
pub fn reserve_order(&mut self) -> u64 {
let order = self.next_order;
self.next_order = self.next_order.saturating_add(1);
order
}
fn record_order(&mut self, acceptance_order: Option<u64>) -> u64 {
let order = acceptance_order.unwrap_or(self.next_order);
self.next_order = self.next_order.max(order.saturating_add(1));
order
}
pub fn verified_answers(&self) -> impl DoubleEndedIterator<Item = &VerifiedAnswer> {
self.verified_answers.iter().map(|entry| &entry.value)
}
@@ -186,8 +208,13 @@ impl RetainedContext {
entries.into_iter().map(|(_, entry)| entry)
}
/// Records the already-persisted user item; checkpoint/suffix replay uses this same path.
pub fn record_user_message(&mut self, mut message: RetainedUserMessage) {
/// Records a delivered user item with its acceptance order. Legacy items without
/// this metadata use recording order; checkpoint/suffix replay uses the same path.
pub fn record_user_message(
&mut self,
mut message: RetainedUserMessage,
acceptance_order: Option<u64>,
) {
message.bound();
if let Some(index) = self.user_messages.iter().position(|entry| {
message.message_id.is_some() && entry.value.message_id == message.message_id
@@ -197,11 +224,11 @@ impl RetainedContext {
}
self.user_messages.remove(index);
}
let order = self.record_order(acceptance_order);
self.user_messages.push_back(Ordered {
order: self.next_order,
order,
value: message,
});
self.next_order = self.next_order.saturating_add(1);
bound_family(&mut self.user_messages, &mut self.user_messages_incomplete);
}
@@ -210,7 +237,10 @@ impl RetainedContext {
let mut event = event.clone();
event.bound();
match event {
RetainedContextEvent::VerifiedAnswer(answer) => {
RetainedContextEvent::VerifiedAnswer {
answer,
acceptance_order,
} => {
if let Some(index) = self.verified_answers.iter().position(|existing| {
existing.value.turn_id == answer.turn_id
&& existing.value.call_id == answer.call_id
@@ -220,11 +250,11 @@ impl RetainedContext {
}
self.verified_answers.remove(index);
}
let order = self.record_order(acceptance_order);
self.verified_answers.push_back(Ordered {
order: self.next_order,
order,
value: answer,
});
self.next_order = self.next_order.saturating_add(1);
bound_family(
&mut self.verified_answers,
&mut self.verified_answers_incomplete,
@@ -242,9 +272,12 @@ impl RetainedContext {
..Self::default()
});
for entry in &mut self.verified_answers {
let mut event = RetainedContextEvent::VerifiedAnswer(entry.value.clone());
let mut event = RetainedContextEvent::VerifiedAnswer {
answer: entry.value.clone(),
acceptance_order: Some(entry.order),
};
event.bound();
let RetainedContextEvent::VerifiedAnswer(answer) = event;
let RetainedContextEvent::VerifiedAnswer { answer, .. } = event;
entry.value = answer;
self.next_order = self.next_order.max(entry.order.saturating_add(1));
}
@@ -264,15 +297,23 @@ impl RetainedContext {
self.verified_answers.retain(|answer| keep(&answer.value));
}
/// Rolls back at the original user-message boundary, including facts delivered after it.
/// Rolls back at the original user-message boundary, including later-accepted facts.
/// The explicit order also covers checkpoints made before a queued message was delivered.
/// Steering can share a turn ID. Legacy sources without message identity fall back to
/// source-turn removal and cannot establish complete retained user instructions.
pub fn rollback(&mut self, turn_ids: &[&str], first_removed_message_id: Option<&str>) {
if let Some(order) = first_removed_message_id.and_then(|id| {
self.user_messages
.iter()
.find(|message| message.value.message_id.as_deref() == Some(id))
.map(|message| message.order)
pub fn rollback(
&mut self,
turn_ids: &[&str],
first_removed_message_id: Option<&str>,
acceptance_order: Option<u64>,
) {
if let Some(order) = acceptance_order.or_else(|| {
first_removed_message_id.and_then(|id| {
self.user_messages
.iter()
.find(|message| message.value.message_id.as_deref() == Some(id))
.map(|message| message.order)
})
}) {
self.verified_answers.retain(|entry| entry.order < order);
self.user_messages.retain(|entry| entry.order < order);
codex-rs/history/src/retained_context_tests.rs
+125 -38
View File
@@ -2,14 +2,17 @@ use super::*;
use pretty_assertions::assert_eq;
fn publish_answer() -> RetainedContextEvent {
RetainedContextEvent::VerifiedAnswer(VerifiedAnswer {
turn_id: "turn-1".to_owned(),
call_id: "ask-1".to_owned(),
questions: vec![VerifiedQuestionAnswer {
question: "Publish?".to_owned(),
answer: "Yes, but never publicly.".to_owned(),
}],
})
RetainedContextEvent::VerifiedAnswer {
answer: VerifiedAnswer {
turn_id: "turn-1".to_owned(),
call_id: "ask-1".to_owned(),
questions: vec![VerifiedQuestionAnswer {
question: "Publish?".to_owned(),
answer: "Yes, but never publicly.".to_owned(),
}],
},
acceptance_order: None,
}
}
#[test]
@@ -18,12 +21,15 @@ fn retained_evidence_preserves_order_through_checkpoint_and_rollback() {
let first = publish_answer();
assert!(context.record(&first));
let before_restriction = context.clone();
context.record_user_message(RetainedUserMessage {
turn_id: "revocation-turn".to_owned(),
message_id: Some("revocation".to_owned()),
text: "Do not publish after all.".to_owned(),
complete: true,
});
context.record_user_message(
RetainedUserMessage {
turn_id: "revocation-turn".to_owned(),
message_id: Some("revocation".to_owned()),
text: "Do not publish after all.".to_owned(),
complete: true,
},
/*acceptance_order*/ None,
);
let snapshot = context.clone();
assert!(!context.record(&first));
assert_eq!(context, snapshot);
@@ -43,7 +49,11 @@ fn retained_evidence_preserves_order_through_checkpoint_and_rollback() {
.collect::<Vec<_>>(),
vec!["Yes, but never publicly.", "Do not publish after all."]
);
restored.rollback(&["revocation-turn"], Some("revocation"));
restored.rollback(
&["revocation-turn"],
Some("revocation"),
/*acceptance_order*/ None,
);
assert_eq!(
restored,
RetainedContext {
@@ -61,24 +71,31 @@ fn retained_families_enforce_storage_limits_without_changing_snapshots() {
let snapshot = context.clone();
for index in 2..=10 {
context.record(&RetainedContextEvent::VerifiedAnswer(VerifiedAnswer {
turn_id: "turn-2".to_owned(),
call_id: format!("ask-{index}"),
questions: vec![VerifiedQuestionAnswer {
question: "Continue?".to_owned(),
answer: "Yes".to_owned(),
}],
}));
context.record(&RetainedContextEvent::VerifiedAnswer {
answer: VerifiedAnswer {
turn_id: "turn-2".to_owned(),
call_id: format!("ask-{index}"),
questions: vec![VerifiedQuestionAnswer {
question: "Continue?".to_owned(),
answer: "Yes".to_owned(),
}],
},
acceptance_order: None,
});
}
assert!(!context.verified_answers_complete());
assert_eq!(context.verified_answers().count(), MAX_FAMILY_RECORDS);
context.rollback(&["turn-2"], /*first_removed_message_id*/ None);
context.rollback(
&["turn-2"],
/*first_removed_message_id*/ None,
/*acceptance_order*/ None,
);
assert_eq!(context.verified_answers().count(), 0);
assert!(!context.verified_answers_complete());
let mut restored = snapshot.clone();
let mut oversized = first.clone();
let RetainedContextEvent::VerifiedAnswer(answer) = &mut oversized;
let RetainedContextEvent::VerifiedAnswer { answer, .. } = &mut oversized;
answer.questions[0].answer = "a".repeat(MAX_RECORD_BYTES);
restored.record(&oversized);
assert!(!restored.verified_answers_complete());
@@ -100,12 +117,15 @@ fn retained_families_enforce_storage_limits_without_changing_snapshots() {
"Yes, but never publicly."
);
for index in 0..=MAX_FAMILY_RECORDS {
restored.record_user_message(RetainedUserMessage {
turn_id: "later-turn".to_owned(),
message_id: Some(format!("message-{index}")),
text: "Keep the repository private.".to_owned(),
complete: true,
});
restored.record_user_message(
RetainedUserMessage {
turn_id: "later-turn".to_owned(),
message_id: Some(format!("message-{index}")),
text: "Keep the repository private.".to_owned(),
complete: true,
},
/*acceptance_order*/ None,
);
}
assert!(!restored.user_messages_complete());
assert_eq!(
@@ -115,12 +135,29 @@ fn retained_families_enforce_storage_limits_without_changing_snapshots() {
.count(),
MAX_FAMILY_RECORDS
);
restored.record_user_message(RetainedUserMessage {
turn_id: "oversized-turn".to_owned(),
message_id: Some("oversized-message".to_owned()),
text: "restriction ".repeat(MAX_RECORD_BYTES),
complete: true,
});
let recent = restored.clone();
restored.record_user_message(
RetainedUserMessage {
turn_id: "earlier-turn".to_owned(),
message_id: Some("delayed-message".to_owned()),
text: "An older queued instruction.".to_owned(),
complete: true,
},
Some(0),
);
assert_eq!(
restored, recent,
"delayed old input must not evict newer evidence"
);
restored.record_user_message(
RetainedUserMessage {
turn_id: "oversized-turn".to_owned(),
message_id: Some("oversized-message".to_owned()),
text: "restriction ".repeat(MAX_RECORD_BYTES),
complete: true,
},
/*acceptance_order*/ None,
);
let Some(RetainedContextEntry::UserMessage(message)) = restored.ordered_entries().next_back()
else {
panic!("latest user evidence");
@@ -130,7 +167,7 @@ fn retained_families_enforce_storage_limits_without_changing_snapshots() {
#[test]
fn legacy_checkpoints_mark_user_messages_incomplete() {
let RetainedContextEvent::VerifiedAnswer(answer) = publish_answer();
let RetainedContextEvent::VerifiedAnswer { answer, .. } = publish_answer();
let mut wire = serde_json::json!({
"verified_answers": [answer], "incomplete": false
});
@@ -150,3 +187,53 @@ fn legacy_checkpoints_mark_user_messages_incomplete() {
restored.restore(/*checkpoint*/ None);
assert!(!restored.user_messages_complete());
}
#[test]
fn accepted_order_survives_delayed_recording_and_checkpoint_replay() {
let mut context = RetainedContext::default();
// Rejected/canceled input consumes a sequence but never becomes evidence.
context.reserve_order();
let steer_order = context.reserve_order();
let answer_order = context.reserve_order();
let RetainedContextEvent::VerifiedAnswer { answer, .. } = publish_answer();
let event = RetainedContextEvent::VerifiedAnswer {
answer,
acceptance_order: Some(answer_order),
};
context.record(&event);
let checkpoint = context.clone();
let instruction = RetainedUserMessage {
turn_id: "turn-1".to_owned(),
message_id: Some("steer".to_owned()),
text: "Keep the repository private.".to_owned(),
complete: true,
};
context.record_user_message(instruction.clone(), Some(steer_order));
let mut resumed = RetainedContext::default();
resumed.restore(Some(&checkpoint));
resumed.record_user_message(instruction, Some(steer_order));
assert_eq!(resumed, context);
assert_eq!(
resumed
.ordered_entries()
.map(|entry| match entry {
RetainedContextEntry::UserMessage(message) => message.text.as_str(),
RetainedContextEntry::VerifiedAnswer(answer) => answer.questions[0].answer.as_str(),
})
.collect::<Vec<_>>(),
vec!["Keep the repository private.", "Yes, but never publicly."],
);
// This checkpoint predates model delivery of the queued instruction. Rollback
// still removes its later-accepted answer using the persisted boundary order.
let mut before_delivery = checkpoint;
before_delivery.rollback(&["turn-1"], Some("steer"), Some(steer_order));
resumed.rollback(&["turn-1"], Some("steer"), Some(steer_order));
assert_eq!(before_delivery, resumed);
assert_eq!(
resumed,
RetainedContext {
next_order: 3,
..Default::default()
}
);
}
codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs
+46 -18
View File
@@ -42,9 +42,16 @@ struct PendingUserResponse {
struct InstructionBoundary {
record_index: usize,
message_id: Option<ResponseItemId>,
acceptance_order: Option<u64>,
alive: bool,
}
struct RetainedFactSource {
record_index: usize,
turn_id: String,
acceptance_order: Option<u64>,
}
/// Compact plan keyed by parsed source-record index.
pub(super) struct RollbackPlan {
record_boundaries: Vec<Option<usize>>,
@@ -97,7 +104,7 @@ pub(super) struct RollbackPlanner {
pending_delivery_boundary: Option<usize>,
turn_boundaries: HashMap<String, usize>,
call_boundaries: HashMap<(String, String), Option<usize>>,
retained_fact_sources: Vec<(usize, String)>,
retained_fact_sources: Vec<RetainedFactSource>,
compactions: Vec<CompactionFrame>,
model_replay: ModelReplayPlanner,
}
@@ -154,6 +161,10 @@ impl RollbackPlanner {
} else if rollback::counts_as_boundary(&response.item) {
let boundary = self.start_boundary(index);
self.boundaries[boundary].message_id = response.id().cloned();
self.boundaries[boundary].acceptance_order = response
.metadata
.as_ref()
.and_then(|metadata| metadata.user_input_order);
if let ResponseItem::Message { role, content, .. } = &response.item
&& role == "user"
{
@@ -247,9 +258,10 @@ impl RollbackPlanner {
self.assign_targeted_record(index, Some(record.turn_id.as_str()));
}
RolloutItem::WorldState(_) | RolloutItem::RealtimeItem(_) => {}
RolloutItem::RetainedContext(codex_rollout::RetainedContextEvent::VerifiedAnswer(
RolloutItem::RetainedContext(codex_rollout::RetainedContextEvent::VerifiedAnswer {
answer,
)) => {
acceptance_order,
}) => {
let source = (answer.turn_id.clone(), answer.call_id.clone());
// A late answer still belongs to its call's instruction boundary, even
// if the same running turn has since received another user steer.
@@ -260,8 +272,11 @@ impl RollbackPlanner {
}
self.call_boundaries
.insert(source, self.record_boundaries[index]);
self.retained_fact_sources
.push((index, answer.turn_id.clone()));
self.retained_fact_sources.push(RetainedFactSource {
record_index: index,
turn_id: answer.turn_id.clone(),
acceptance_order: *acceptance_order,
});
}
RolloutItem::SecurityRiskScore(_) => self.record_boundaries[index] = None,
}
@@ -308,6 +323,7 @@ impl RollbackPlanner {
self.boundaries.push(InstructionBoundary {
record_index: index,
message_id: None,
acceptance_order: None,
alive: true,
});
let had_prior_boundary = !self.boundary_stack.is_empty();
@@ -364,6 +380,18 @@ impl RollbackPlanner {
}
if let Some(boundary) = first_removed_boundary {
let source = &self.boundaries[boundary];
if let Some(order) = source.acceptance_order {
// An answer may have been persisted before the queued instruction
// accepted ahead of it. Both are removed at that acceptance boundary.
for fact in &self.retained_fact_sources {
if fact
.acceptance_order
.is_some_and(|accepted| accepted >= order)
{
self.record_boundaries[fact.record_index] = Some(boundary);
}
}
}
let removed_turns = self
.turn_boundaries
.iter()
@@ -372,29 +400,29 @@ impl RollbackPlanner {
.chain(
self.retained_fact_sources
.iter()
.filter(|(index, _)| {
self.record_boundaries[*index]
.filter(|fact| {
self.record_boundaries[fact.record_index]
.is_some_and(|boundary| removed_boundaries.contains(&boundary))
})
.map(|(_, turn_id)| turn_id.as_str()),
.map(|fact| fact.turn_id.as_str()),
)
.collect::<Vec<_>>();
// A checkpoint before this instruction cannot contain evidence from it.
// Accepted answers can precede a queued instruction in the rollout,
// including in checkpoints. Legacy evidence uses the recorded boundary.
// Later checkpoints have not been observed yet and already reflect this rollback.
for frame in self
.compactions
.iter_mut()
.rev()
.take_while(|frame| frame.record_index >= source.record_index)
{
for frame in self.compactions.iter_mut().rev().take_while(|frame| {
source.acceptance_order.is_some() || frame.record_index >= source.record_index
}) {
if let Some(context) = &mut frame.item.retained_context {
if context
.ordered_entries()
.any(|entry| matches!(entry, RetainedContextEntry::UserMessage(_)))
if source.acceptance_order.is_some()
|| context
.ordered_entries()
.any(|entry| matches!(entry, RetainedContextEntry::UserMessage(_)))
{
context.rollback(
&removed_turns,
source.message_id.as_ref().map(ResponseItemId::as_str),
source.acceptance_order,
);
} else {
// Checkpoints written without instruction retention keep the legacy
codex-rs/thread-store/src/local/rollout_migration_tests.rs
+49 -9
View File
@@ -1230,7 +1230,12 @@ async fn migration_preserves_answers_before_a_rolled_back_steer() {
checkpoint
.verified_answers()
.cloned()
.map(codex_rollout::RetainedContextEvent::VerifiedAnswer)
.map(
|answer| codex_rollout::RetainedContextEvent::VerifiedAnswer {
answer,
acceptance_order: None,
}
)
.collect::<Vec<_>>(),
answers[..1]
);
@@ -1272,6 +1277,15 @@ async fn migration_preserves_reverse_replay_anchor_after_pre_compaction_rollback
#[tokio::test]
async fn migration_preserves_same_turn_evidence_across_rollback_checkpoints() {
assert_migrated_evidence_order(/*steer_order*/ None).await;
}
#[tokio::test]
async fn migration_removes_answers_accepted_after_a_queued_steer() {
assert_migrated_evidence_order(Some(1)).await;
}
async fn assert_migrated_evidence_order(steer_order: Option<u64>) {
const INITIAL: &str = "Never publish publicly.";
const STEER: &str = "Also inspect the README.";
let home = TempDir::new().expect("create Codex home");
@@ -1299,10 +1313,10 @@ async fn migration_preserves_same_turn_evidence_across_rollback_checkpoints() {
let retained = json!({
"user_messages": [
{"order": 0, "turn_id": "shared-turn", "message_id": "msg_initial", "text": INITIAL, "complete": true},
{"order": 2, "turn_id": "shared-turn", "message_id": "msg_steer", "text": STEER, "complete": true}
{"order": steer_order.unwrap_or(2), "turn_id": "shared-turn", "message_id": "msg_steer", "text": STEER, "complete": true}
],
"verified_answers": [
{"order": 1, "turn_id": "shared-turn", "call_id": "before", "questions": [{"question": "Publish?", "answer": "Only privately."}]},
{"order": if steer_order.is_some() { 2 } else { 1 }, "turn_id": "shared-turn", "call_id": "before", "questions": [{"question": "Publish?", "answer": "Only privately."}]},
{"order": 3, "turn_id": "shared-turn", "call_id": "after", "questions": [{"question": "Publish the README?", "answer": "Do not publish it."}]}
],
"incomplete": false, "user_messages_incomplete": false, "next_order": 4
@@ -1315,8 +1329,26 @@ async fn migration_preserves_same_turn_evidence_across_rollback_checkpoints() {
.expect("retained checkpoint")
.verified_answers()
.cloned()
.map(codex_rollout::RetainedContextEvent::VerifiedAnswer)
.map(|answer| {
let acceptance_order =
steer_order.map(|_| if answer.call_id == "before" { 2 } else { 3 });
codex_rollout::RetainedContextEvent::VerifiedAnswer {
answer,
acceptance_order,
}
})
.collect::<Vec<_>>();
let mut before_retained = retained.clone();
before_retained["user_messages"]
.as_array_mut()
.expect("user messages")
.truncate(/*len*/ 1);
before_retained["verified_answers"]
.as_array_mut()
.expect("verified answers")
.truncate(/*len*/ 1);
before_retained["next_order"] = json!(if steer_order.is_some() { 3 } else { 2 });
let remaining_answers = usize::from(steer_order.is_none());
let mut expected = retained;
expected["user_messages"]
.as_array_mut()
@@ -1325,13 +1357,13 @@ async fn migration_preserves_same_turn_evidence_across_rollback_checkpoints() {
expected["verified_answers"]
.as_array_mut()
.expect("verified answers")
.truncate(/*len*/ 1);
.truncate(remaining_answers);
let mut before_expected = expected.clone();
before_expected["next_order"] = json!(2);
before_expected["next_order"] = before_retained["next_order"].clone();
let mut before_steer = checkpoint.clone();
before_steer.replacement_history = Some(vec![initial.clone().into()]);
before_steer.retained_context =
Some(serde_json::from_value(before_expected.clone()).expect("pre-steer checkpoint"));
Some(serde_json::from_value(before_retained).expect("pre-steer checkpoint"));
let mut after_rollback = before_steer.clone();
after_rollback.retained_context =
Some(serde_json::from_value(expected.clone()).expect("post-rollback checkpoint"));
@@ -1345,7 +1377,15 @@ async fn migration_preserves_same_turn_evidence_across_rollback_checkpoints() {
user_message(INITIAL),
RolloutItem::RetainedContext(answers[0].clone()),
RolloutItem::Compacted(before_steer),
rollout_response_item(steer),
RolloutItem::ResponseItem(codex_rollout::ResponseItemEnvelope {
item: steer,
metadata: steer_order.map(|order| {
serde_json::from_value(json!({
"user_input_order": order
}))
.expect("acceptance metadata")
}),
}),
user_message(STEER),
RolloutItem::RetainedContext(answers[1].clone()),
completed("shared-turn"),
@@ -1399,7 +1439,7 @@ async fn migration_preserves_same_turn_evidence_across_rollback_checkpoints() {
_ => None,
})
.collect::<Vec<_>>(),
vec![answers[0].clone()]
answers[..remaining_answers]
);
}