diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index 3b9f9d95f1..bd13ac46e8 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 19ec09fab6..ab8cf48895 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -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::>(); 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. diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index d23e8ef8ad..98ad19ef8e 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -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()); } } diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 4964905418..e17f1508e7 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -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; diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index 1d4f8fc26c..6a9898c562 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -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(), }), }; diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 9d44f614ca..ca80687ffc 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -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, ) } diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index 3993abf7dc..94d3dfeb4d 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -23,6 +23,8 @@ pub enum TurnInput { UserInput { content: Vec, client_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + acceptance_order: Option, }, 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(), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 00d3e8e252..1e05614090 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -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 { + ) -> Option { 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, + acceptance_order: Option, 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); diff --git a/codex-rs/core/src/session/retained_context.rs b/codex-rs/core/src/session/retained_context.rs index 9fa962856b..71506e68f4 100644 --- a/codex-rs/core/src/session/retained_context.rs +++ b/codex-rs/core/src/session/retained_context.rs @@ -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 { + 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 diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index a2453d227b..202aae4aca 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -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. diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index ef0ce0503d..c76b9b4fec 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -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); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d9c1f72214..f302a92a07 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -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(), diff --git a/codex-rs/core/src/session/turn_input.rs b/codex-rs/core/src/session/turn_input.rs index 79df5facfa..fe7b5ce32b 100644 --- a/codex-rs/core/src/session/turn_input.rs +++ b/codex-rs/core/src/session/turn_input.rs @@ -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, diff --git a/codex-rs/core/src/session/turn_input_tests.rs b/codex-rs/core/src/session/turn_input_tests.rs index c9810df1f4..8c2d7c5b6a 100644 --- a/codex-rs/core/src/session/turn_input_tests.rs +++ b/codex-rs/core/src/session/turn_input_tests.rs @@ -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(), diff --git a/codex-rs/core/src/state/mod.rs b/codex-rs/core/src/state/mod.rs index 16297ea417..218b25ba92 100644 --- a/codex-rs/core/src/state/mod.rs +++ b/codex-rs/core/src/state/mod.rs @@ -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; diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index c61e6431e9..b03dbf737f 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -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>, pending_request_permissions: HashMap, - pending_user_input: HashMap>, + pending_user_input: HashMap>, pending_elicitations: HashMap<(String, RequestId), oneshot::Sender>, mcp_tool_approval_metadata: HashMap, McpToolApprovalMetadata)>, pending_dynamic_tools: HashMap>, @@ -106,6 +106,13 @@ pub(crate) struct TurnState { pub(crate) last_known_step_context: Option>, } +/// 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, +} + pub(crate) struct PendingRequestPermissions { pub(crate) tx_response: oneshot::Sender, 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, - ) -> Option> { + tx: oneshot::Sender, + ) -> Option> { self.pending_user_input.insert(key, tx) } pub(crate) fn remove_pending_user_input( &mut self, key: &str, - ) -> Option> { + ) -> Option> { self.pending_user_input.remove(key) } diff --git a/codex-rs/core/src/tools/handlers/request_user_input.rs b/codex-rs/core/src/tools/handlers/request_user_input.rs index 10562985e4..a00781ed24 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input.rs @@ -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::>(); 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; } } diff --git a/codex-rs/core/tests/suite/guardian_retained_context.rs b/codex-rs/core/tests/suite/guardian_retained_context.rs index 0d2048fe3e..3c0c7b9442 100644 --- a/codex-rs/core/tests/suite/guardian_retained_context.rs +++ b/codex-rs/core/tests/suite/guardian_retained_context.rs @@ -51,6 +51,7 @@ async fn record_answer( server: &MockServer, call_id: &str, answer: &str, + acceptance_order: u64, ) -> Result { 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::>(); @@ -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::>(), ) @@ -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]; diff --git a/codex-rs/core/tests/suite/hooks.rs b/codex-rs/core/tests/suite/hooks.rs index 090ab49a0a..112f64e08d 100644 --- a/codex-rs/core/tests/suite/hooks.rs +++ b/codex-rs/core/tests/suite/hooks.rs @@ -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::>(), + 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!( diff --git a/codex-rs/history/src/lib.rs b/codex-rs/history/src/lib.rs index a77b1a2e66..43d0231ee2 100644 --- a/codex-rs/history/src/lib.rs +++ b/codex-rs/history/src/lib.rs @@ -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, + + /// 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, } impl ResponseItemEnvelope { diff --git a/codex-rs/history/src/retained_context.rs b/codex-rs/history/src/retained_context.rs index 350dbdf84b..f02987367d 100644 --- a/codex-rs/history/src/retained_context.rs +++ b/codex-rs/history/src/retained_context.rs @@ -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, + }, } /// 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>, @@ -98,6 +104,9 @@ fn legacy_user_messages_incomplete() -> bool { } fn bound_family(items: &mut VecDeque>, 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 { + 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 { 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, + ) { 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, + ) { + 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); diff --git a/codex-rs/history/src/retained_context_tests.rs b/codex-rs/history/src/retained_context_tests.rs index 0e945c47e2..2157ee3022 100644 --- a/codex-rs/history/src/retained_context_tests.rs +++ b/codex-rs/history/src/retained_context_tests.rs @@ -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!["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!["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() + } + ); +} diff --git a/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs b/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs index b16539ea7e..38f04e8c7a 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs +++ b/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs @@ -42,9 +42,16 @@ struct PendingUserResponse { struct InstructionBoundary { record_index: usize, message_id: Option, + acceptance_order: Option, alive: bool, } +struct RetainedFactSource { + record_index: usize, + turn_id: String, + acceptance_order: Option, +} + /// Compact plan keyed by parsed source-record index. pub(super) struct RollbackPlan { record_boundaries: Vec>, @@ -97,7 +104,7 @@ pub(super) struct RollbackPlanner { pending_delivery_boundary: Option, turn_boundaries: HashMap, call_boundaries: HashMap<(String, String), Option>, - retained_fact_sources: Vec<(usize, String)>, + retained_fact_sources: Vec, compactions: Vec, 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::>(); - // 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 diff --git a/codex-rs/thread-store/src/local/rollout_migration_tests.rs b/codex-rs/thread-store/src/local/rollout_migration_tests.rs index 0758db8163..82339f04dc 100644 --- a/codex-rs/thread-store/src/local/rollout_migration_tests.rs +++ b/codex-rs/thread-store/src/local/rollout_migration_tests.rs @@ -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::>(), 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) { 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::>(); + 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![answers[0].clone()] + answers[..remaining_answers] ); }