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 520cb64c9c..3b9f9d95f1 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/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 3b7ca36c58..e928f5290e 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -2101,10 +2101,9 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() { ), "a subagent must not inherit its parent review checkpoint", ); - assert_eq!( - history.retained_context(), - &codex_history::RetainedContext::default() - ); + let mut expected_retained_context = codex_history::RetainedContext::default(); + expected_retained_context.mark_user_messages_incomplete(); + assert_eq!(history.retained_context(), &expected_retained_context); assert!( history_contains_text(history.raw_items(), "compacted parent summary"), "forked child history should retain compacted non-hint content" diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 5e473b0547..796ae1bb22 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -32,8 +32,10 @@ use codex_analytics::CompactionTrigger; use codex_analytics::now_unix_seconds; use codex_context_fragments::AnnotatedContent; use codex_context_fragments::set_annotated_content; +use codex_features::Feature; use codex_history::CodexHarnessMetadata; use codex_history::ResponseItemEnvelope; +use codex_protocol::ResponseItemId; use codex_protocol::error::CodexErr; use codex_protocol::error::CodexErrorDetails; use codex_protocol::error::Result as CodexResult; @@ -355,7 +357,12 @@ async fn run_compact_task_inner_impl( let summary_suffix = get_last_assistant_message_from_turn(history_snapshot.raw_items()).unwrap_or_default(); let summary_text = format!("{SUMMARY_PREFIX}\n{summary_suffix}"); - let user_messages = collect_annotated_user_messages(history_items); + let identity = if sess.enabled(Feature::GuardianThreadContext) { + CompactedMessageIdentity::Preserve + } else { + CompactedMessageIdentity::Regenerate + }; + let user_messages = collect_annotated_user_messages(history_items, identity); let mut new_history = build_compacted_history(Vec::new(), &user_messages, &summary_text); if let Some(summary_item) = new_history.last_mut() { @@ -527,6 +534,9 @@ pub fn content_items_to_text(content: &[ContentItem]) -> Option { #[derive(Clone, Debug, PartialEq)] pub(crate) struct CompactedUserMessage { + // Keep source identity even when compaction shortens the text, so rollback can + // correlate the rebuilt message with thread-owned retained evidence. + id: Option, message: String, internal_chat_message_metadata_passthrough: Option, harness_metadata: Option, @@ -540,12 +550,24 @@ pub(crate) fn collect_user_messages(items: &[ResponseItem]) -> Vec Vec { items .iter() .filter_map(|envelope| compacted_user_message(&envelope.item, envelope.metadata.clone())) + .map(|mut message| { + if matches!(identity, CompactedMessageIdentity::Regenerate) { + message.id = None; + } + message + }) .collect() } @@ -560,6 +582,7 @@ fn compacted_user_message( return None; } Some(CompactedUserMessage { + id: item.id().cloned(), message: user.message(), internal_chat_message_metadata_passthrough: match item { ResponseItem::Message { @@ -678,6 +701,7 @@ fn build_compacted_history_with_limit( let truncated = truncate_text(&message.message, TruncationPolicy::Tokens(remaining)); selected_messages.push(CompactedUserMessage { + id: message.id.clone(), message: truncated, internal_chat_message_metadata_passthrough: message .internal_chat_message_metadata_passthrough @@ -692,7 +716,7 @@ fn build_compacted_history_with_limit( for message in &selected_messages { let mut item = ResponseItem::Message { - id: None, + id: message.id.clone(), role: "user".to_string(), content: vec![ContentItem::InputText { text: message.message.clone(), diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index 44cfad1995..72fa09a1e8 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -66,6 +66,7 @@ fn user_message(text: &str) -> ResponseItem { fn compacted_user_message(text: &str) -> CompactedUserMessage { CompactedUserMessage { + id: None, message: text.to_string(), internal_chat_message_metadata_passthrough: None, harness_metadata: None, @@ -129,7 +130,13 @@ fn collect_user_messages_extracts_user_text_only() { let collected = collect_user_messages(&items); - assert_eq!(vec![compacted_user_message("first")], collected); + assert_eq!( + vec![CompactedUserMessage { + id: Some(ResponseItemId::with_suffix("msg", "user")), + ..compacted_user_message("first") + }], + collected, + ); } #[test] @@ -142,10 +149,11 @@ fn collect_annotated_user_messages_extracts_user_text_only() { ResponseItemEnvelope::new(ResponseItem::Other), ]; - let collected = collect_annotated_user_messages(&items); + let collected = collect_annotated_user_messages(&items, CompactedMessageIdentity::Preserve); assert_eq!( vec![CompactedUserMessage { + id: None, message: "first".to_string(), internal_chat_message_metadata_passthrough: None, harness_metadata: Some(CodexHarnessMetadata::default()), @@ -223,6 +231,7 @@ fn build_token_limited_compacted_history_truncates_overlong_user_messages() { let max_tokens = 16; let big = "word ".repeat(200); let user_message = CompactedUserMessage { + id: Some(ResponseItemId::with_suffix("msg", "long-user")), message: big.clone(), internal_chat_message_metadata_passthrough: None, harness_metadata: Some(CodexHarnessMetadata::default()), @@ -261,6 +270,7 @@ fn build_token_limited_compacted_history_truncates_overlong_user_messages() { other => panic!("unexpected item in history: {other:?}"), }; assert_eq!(summary_text, "SUMMARY"); + assert_eq!(history[0].id(), user_message.id.as_ref()); assert_eq!(history[0].metadata, Some(CodexHarnessMetadata::default())); assert_eq!(history[1].metadata, None); } @@ -292,6 +302,7 @@ fn build_compacted_history_preserves_user_message_passthrough_metadata() { let history = build_compacted_history( Vec::new(), &[CompactedUserMessage { + id: Some(ResponseItemId::with_suffix("msg", "user")), message: "first user message".to_string(), internal_chat_message_metadata_passthrough: Some( InternalChatMessageMetadataPassthrough { @@ -314,7 +325,7 @@ fn build_compacted_history_preserves_user_message_passthrough_metadata() { vec![ ResponseItemEnvelope { item: ResponseItem::Message { - id: None, + id: Some(ResponseItemId::with_suffix("msg", "user")), role: "user".to_string(), content: vec![ContentItem::InputText { text: "first user message".to_string(), diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index b02d2188d8..19ec09fab6 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -61,6 +61,8 @@ pub(crate) struct ContextManager { review_history: Option, /// Host facts independent of the model window; snapshots share immutable state. retained_context: Arc, + /// Live and replay instruction capture are enabled together by the session feature flag. + retain_user_messages: bool, /// Bumped whenever history is rewritten, such as compaction or rollback. history_version: u64, /// Monotonic user-input/reset revision, independent of compaction's history generation. @@ -156,6 +158,7 @@ impl ContextManager { items: Arc::new(Vec::new()), review_history: None, retained_context: Arc::default(), + retain_user_messages: false, history_version: 0, user_message_revision: 0, token_info: TokenUsageInfo::new_or_append( @@ -180,6 +183,10 @@ impl ContextManager { &self.retained_context } + pub(crate) fn enable_user_message_retention(&mut self) { + self.retain_user_messages = true; + } + pub(crate) fn record_retained_context(&mut self, event: &RetainedContextEvent) -> bool { if !Arc::make_mut(&mut self.retained_context).record(event) { return false; @@ -188,7 +195,7 @@ impl ContextManager { true } - pub(crate) fn restore_retained_context(&mut self, checkpoint: &RetainedContext) { + pub(crate) fn restore_retained_context(&mut self, checkpoint: Option<&RetainedContext>) { Arc::make_mut(&mut self.retained_context).restore(checkpoint); } @@ -318,6 +325,44 @@ impl ContextManager { } Arc::make_mut(&mut self.items).push(processed); if crate::context::is_user_authorization_message(item) { + if self.retain_user_messages + && let ResponseItem::Message { + content, + internal_chat_message_metadata_passthrough, + .. + } = item + { + let mut complete = internal_chat_message_metadata_passthrough + .as_ref() + .and_then(|metadata| metadata.content_item_kinds.as_ref()) + .is_some_and(|kinds| { + kinds.len() == content.len() + && kinds.iter().all(|kind| kind.0.starts_with("user.")) + }); + let text = content + .iter() + .filter_map(|content| match content { + ContentItem::InputText { text } | ContentItem::OutputText { text } => { + Some(text.as_str()) + } + _ => { + complete = false; + None + } + }) + .collect::>() + .join("\n"); + Arc::make_mut(&mut self.retained_context).record_user_message( + codex_history::RetainedUserMessage { + turn_id: item.turn_id().unwrap_or_default().to_owned(), + message_id: item.id().map(|id| id.as_str().to_owned()), + text, + complete, + }, + ); + } else { + Arc::make_mut(&mut self.retained_context).mark_user_messages_incomplete(); + } self.user_message_revision = self.user_message_revision.saturating_add(1); } } @@ -488,6 +533,9 @@ impl ContextManager { user_positions[user_positions.len() - n_from_end] }; + let first_removed_message_id = snapshot[cut_idx] + .id() + .map(codex_protocol::ResponseItemId::as_str); let mut review_history = self.review_history.take(); if let Some(history) = &mut review_history { history.truncate_before(&snapshot[cut_idx].item); @@ -528,18 +576,21 @@ impl ContextManager { .iter() .filter_map(|item| item.turn_id()) .collect::>(); - Arc::make_mut(&mut retained_context).retain_answers(|answer| { - // A steer creates an instruction boundary, not a new turn ID. Replay exposes - // the original calls from rolled-back checkpoints, so prefer the exact source. - if let Some(source_index) = snapshot.iter().rposition(|item| { - item.turn_id() == Some(answer.turn_id.as_str()) - && matches!(&item.item, ResponseItem::FunctionCall { call_id, .. } - if call_id == &answer.call_id) - }) { - return source_index < cut_idx; - } - !removed_turns.contains(&answer.turn_id.as_str()) - }); + if self.retain_user_messages { + Arc::make_mut(&mut retained_context).rollback(&removed_turns, first_removed_message_id); + } else { + Arc::make_mut(&mut retained_context).retain_answers(|answer| { + // Legacy answers follow their original call, not later steers in the same turn. + if let Some(source_index) = snapshot.iter().rposition(|item| { + item.turn_id() == Some(answer.turn_id.as_str()) + && matches!(&item.item, ResponseItem::FunctionCall { call_id, .. } + if call_id == &answer.call_id) + }) { + return source_index < cut_idx; + } + !removed_turns.contains(&answer.turn_id.as_str()) + }); + } self.replace_annotated(retained_items); self.retained_context = retained_context; self.review_history = review_history; diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index c6e926ee0a..d23e8ef8ad 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -1265,6 +1265,55 @@ fn drop_last_n_user_turns_preserves_prefix() { let retained = history.retained_context().clone(); history.drop_last_n_user_turns(/*num_turns*/ 1); assert_eq!(history.retained_context(), &retained); + + // A steered message shares its source turn, but rollback must keep the earlier + // instruction and answer as complete evidence, including after the next compaction. + let mut history = ContextManager::default(); + history.enable_user_message_retention(); + let mut expected = None; + for (id, text) in [ + ("restriction", "Never publish publicly."), + ("steer", "Check the tests too."), + ] { + let message = ResponseItem::Message { + id: Some(ResponseItemId::with_suffix("msg", id)), + role: "user".to_owned(), + content: vec![ContentItem::InputText { + text: text.to_owned(), + }], + phase: None, + internal_chat_message_metadata_passthrough: Some( + InternalChatMessageMetadataPassthrough { + turn_id: Some("shared-turn".to_owned()), + content_item_kinds: Some(vec![ContentItemKind("user.text".to_owned())]), + ..Default::default() + }, + ), + }; + history.record_items([&message], TruncationPolicy::Tokens(10_000)); + if id == "restriction" { + history.record_retained_context(&codex_history::RetainedContextEvent::VerifiedAnswer( + codex_history::VerifiedAnswer { + turn_id: "shared-turn".to_owned(), + call_id: "ask".to_owned(), + questions: vec![codex_history::VerifiedQuestionAnswer { + question: "Publish?".to_owned(), + answer: "Only privately.".to_owned(), + }], + }, + )); + expected = Some(history.retained_context().clone()); + } + } + history.drop_last_n_user_turns(/*num_turns*/ 1); + history.replace_compacted(Vec::new()); + let retained = history.retained_context(); + assert!(retained.user_messages_complete()); + assert!(retained.verified_answers_complete()); + let mut expected = serde_json::to_value(expected.expect("pre-steer evidence")).unwrap(); + // Rollback removes evidence, but does not reuse its arrival-order sequence numbers. + expected["next_order"] = serde_json::json!(3); + assert_eq!(serde_json::to_value(retained).unwrap(), expected); } #[test] diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d08feebaef..00d3e8e252 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1590,7 +1590,9 @@ impl Session { state .history .restore_guardian_history(guardian_history.as_ref()); - state.history.restore_retained_context(&retained_context); + state + .history + .restore_retained_context(Some(&retained_context)); if let Some(world_state) = world_state_baseline { state.history.set_world_state_baseline(world_state); } diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index a54c08520f..6b5235038c 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -339,15 +339,16 @@ impl Session { .unwrap_or(u64::MAX); let mut history = ContextManager::new(); + if self.enabled(Feature::GuardianThreadContext) { + history.enable_user_message_retention(); + } let mut saw_legacy_compaction_without_replacement_history = false; if let Some(checkpoint) = base_compaction && let Some(items) = &checkpoint.compacted.replacement_history { history.replace_annotated(items.clone()); history.restore_guardian_history(checkpoint.compacted.guardian_history.as_ref()); - if let Some(retained) = &checkpoint.compacted.retained_context { - history.restore_retained_context(retained); - } + history.restore_retained_context(checkpoint.compacted.retained_context.as_ref()); } // Materialize exact history semantics from the replay-derived suffix. The eventual lazy // design should keep this same replay shape, but drive it from a resumable reverse source @@ -386,8 +387,15 @@ impl Session { // prompt shape. // TODO(ccunningham): if we drop support for None replacement_history compaction items, // we can get rid of this second loop entirely and just build `history` directly in the first loop. - let user_messages = - compact::collect_annotated_user_messages(history.annotated_items()); + let identity = if self.enabled(Feature::GuardianThreadContext) { + compact::CompactedMessageIdentity::Preserve + } else { + compact::CompactedMessageIdentity::Regenerate + }; + let user_messages = compact::collect_annotated_user_messages( + history.annotated_items(), + identity, + ); let rebuilt = compact::build_compacted_history( Vec::new(), &user_messages, @@ -395,7 +403,7 @@ impl Session { ); let retained_context = history.retained_context().clone(); history.replace_annotated(rebuilt); - history.restore_retained_context(&retained_context); + history.restore_retained_context(Some(&retained_context)); } } RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => { diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index 93492d9f5c..ef0ce0503d 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -1488,6 +1488,7 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_does_ }], }); let mut retained = codex_history::RetainedContext::default(); + retained.mark_user_messages_incomplete(); retained.record(&answer); let rollout_items = vec![ RolloutItem::ResponseItem(user_message("before compact").into()), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 8d50a6885d..a2d070f547 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1266,6 +1266,9 @@ impl Session { session_configuration.clone(), initial_auto_compact_window_ids, ); + if config.features.enabled(Feature::GuardianThreadContext) { + state.history.enable_user_message_retention(); + } state.base_instructions_provenance = base_instructions_provenance.clone(); let managed_network_requirements_configured = config .config_layer_stack diff --git a/codex-rs/core/tests/suite/guardian_retained_context.rs b/codex-rs/core/tests/suite/guardian_retained_context.rs index 8f45117e48..0d2048fe3e 100644 --- a/codex-rs/core/tests/suite/guardian_retained_context.rs +++ b/codex-rs/core/tests/suite/guardian_retained_context.rs @@ -1,6 +1,6 @@ -//! Retained-answer lifecycle through real sessions and the durable thread-store boundary. -//! Typed fixtures enter via the public rollout append API to isolate storage from producer -//! wiring; compaction, resume, rollback and child forks use production paths. +//! Retained-instruction and answer lifecycles through real sessions and durable checkpoints. +//! Real user input covers steering, compaction and rollback; legacy replay fixtures also use +//! the public rollout append API. Resume and child forks use production paths. use std::collections::HashMap; use std::sync::Arc; @@ -13,11 +13,15 @@ use codex_core::config::ThreadStoreConfig; use codex_features::Feature; use codex_history::InitialHistory; use codex_history::ResumedHistory; +use codex_history::RetainedContext; use codex_history::RetainedContextEvent; use codex_history::RolloutItem; use codex_history::VerifiedAnswer; use codex_history::VerifiedQuestionAnswer; +use codex_protocol::ResponseItemId; use codex_protocol::mcp::ClientMcpExtensions; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::protocol::ThreadHistoryMode; @@ -25,6 +29,7 @@ use codex_protocol::request_user_input::RequestUserInputAnswer; use codex_protocol::request_user_input::RequestUserInputResponse; use codex_protocol::user_input::UserInput; use codex_thread_store::LoadThreadHistoryParams; +use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_function_call_with_namespace; @@ -149,7 +154,7 @@ async fn compact_and_assert_answers( test: &TestCodex, thread: &CodexThread, expected: &[VerifiedAnswer], -) -> Result<()> { +) -> Result { // Inspect live state by persisting a real compaction checkpoint, not a private getter. // Repeating this after resume/rollback also catches checkpoint resurrection. thread.submit(Op::Compact).await?; @@ -172,6 +177,355 @@ async fn compact_and_assert_answers( actual.verified_answers().cloned().collect::>(), expected ); + Ok(actual.clone()) +} + +#[test_case(ThreadHistoryMode::Legacy, true; "enabled legacy rollback")] +#[test_case(ThreadHistoryMode::Paginated, true; "enabled paginated resume")] +#[test_case(ThreadHistoryMode::Legacy, false; "disabled legacy rollback")] +#[test_case(ThreadHistoryMode::Paginated, false; "disabled paginated resume")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn retained_instructions_keep_identity_across_compaction_and_resume( + history_mode: ThreadHistoryMode, + thread_context_enabled: bool, +) -> Result<()> { + skip_if_no_network!(Ok(())); + const INITIAL: &str = "Never publish publicly."; + const STEER: &str = "Also inspect the README."; + // Paginated history supports checkpoint resume, but not the full-history read + // used by ThreadRollback. Exercise rollback on its supported legacy path. + let rollback_counts: &[usize] = match history_mode { + ThreadHistoryMode::Legacy => &[1, 0], + ThreadHistoryMode::Paginated => &[], + }; + let server = start_mock_server().await; + let test = test_codex() + .with_history_mode(history_mode) + .with_config(move |config| { + config.experimental_thread_store = ThreadStoreConfig::Local; + config + .features + .set_enabled(Feature::GuardianThreadContext, thread_context_enabled) + .expect("test context mode"); + // Exercise local compaction's rebuilt user messages, not an opaque checkpoint. + config.model_provider.name = "Local compaction test provider".to_owned(); + config + .features + .disable(Feature::TokenBudget) + .expect("use local compaction"); + config + .features + .enable(Feature::DefaultModeRequestUserInput) + .expect("enable request_user_input"); + }) + .build_with_auto_env(&server) + .await?; + let questions = [ + ("before-steer", "Publish?", "Only privately."), + ("after-steer", "Publish the README?", "Do not publish it."), + ]; + let mut responses = questions + .iter() + .map(|(call_id, question, _)| { + sse(vec![ + ev_function_call( + call_id, + "request_user_input", + &json!({"questions": [{ + "id": "publish", "header": "Publish", "question": question, + "options": [ + {"label": "Yes", "description": "Publish privately."}, + {"label": "No", "description": "Keep local."} + ] + }]}) + .to_string(), + ), + ev_completed(call_id), + ]) + }) + .collect::>(); + responses.push(sse(vec![ev_completed("done")])); + responses.extend((0..=rollback_counts.len()).map(|index| { + sse(vec![ + ev_assistant_message(&format!("summary-{index}"), "Compacted inspection context."), + ev_completed(&format!("compact-{index}")), + ]) + })); + let response_mock = mount_sse_sequence(&server, responses).await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: INITIAL.to_owned(), + text_elements: Vec::new(), + }])) + .await?; + let mut answers = Vec::new(); + for (call_id, question, answer) in questions { + let request = wait_for_event_match(&test.codex, |event| match event { + EventMsg::RequestUserInput(request) => Some(request.clone()), + _ => None, + }) + .await; + if answers.is_empty() { + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: STEER.to_owned(), + text_elements: Vec::new(), + }])) + .await?; + } + test.codex + .submit(Op::UserInputAnswer { + id: request.turn_id.clone(), + response: RequestUserInputResponse { + answers: HashMap::from([( + "publish".to_owned(), + RequestUserInputAnswer { + answers: vec![answer.to_owned()], + }, + )]), + }, + }) + .await?; + answers.push(VerifiedAnswer { + turn_id: request.turn_id, + call_id: call_id.to_owned(), + questions: vec![VerifiedQuestionAnswer { + question: question.to_owned(), + answer: answer.to_owned(), + }], + }); + } + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert_eq!(answers[0].turn_id, answers[1].turn_id); + let requests = response_mock.requests(); + assert_eq!(requests.len(), 3); + assert!(requests[0].has_message_with_input_texts("user", |texts| texts == [INITIAL])); + for (index, request) in requests.iter().skip(1).enumerate() { + assert!(request.has_message_with_input_texts("user", |texts| texts == [STEER])); + for answer in &answers[..=index] { + let (output, _) = request + .function_call_output_content_and_success(&answer.call_id) + .context("answer tool output")?; + let output: serde_json::Value = serde_json::from_str(&output.context("answer text")?)?; + assert_eq!( + output, + json!({"answers": {"publish": {"answers": [answer.questions[0].answer]}}}) + ); + } + } + let history = test.codex.conversation_history_snapshot().await; + let user_messages = [INITIAL, STEER] + .into_iter() + .enumerate() + .map(|(index, text)| { + let message_id = history + .items() + .find(|item| { + matches!( + item, + ResponseItem::Message { role, content, .. } + if role == "user" && matches!( + content.as_slice(), + [ContentItem::InputText { text: body }] if body == text + ) + ) + }) + .and_then(ResponseItem::id) + .expect("original user-message identity"); + json!({ + "order": index * 2, "turn_id": answers[index].turn_id, + "message_id": message_id.as_str(), "text": text, "complete": true, + }) + }) + .collect::>(); + let ordered_answers = answers + .iter() + .enumerate() + .map(|(index, answer)| { + let mut value = json!(answer); + value["order"] = json!(index * 2 + 1); + value + }) + .collect::>(); + let mut expected = json!({ + "user_messages": user_messages, "user_messages_incomplete": false, + "verified_answers": ordered_answers, "incomplete": false, "next_order": 4, + }); + if !thread_context_enabled { + expected = json!({ + "user_messages": [], "user_messages_incomplete": true, + "verified_answers": [], "incomplete": false, "next_order": 0, + }); + answers.clear(); + assert!( + load_context(&test, &test.codex) + .await? + .iter() + .all(|item| !matches!(item, RolloutItem::RetainedContext(_))) + ); + } + 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?)?, + expected + ); + let compacted = test.codex.conversation_history_snapshot().await; + for message in &user_messages { + assert_eq!( + compacted + .items() + .any(|item| item.id().map(ResponseItemId::as_str) == message["message_id"].as_str()), + thread_context_enabled, + "only thread-owned context preserves original user-message identity" + ); + } + let mut thread = Arc::clone(&test.codex); + for &remaining in rollback_counts { + let remaining = if thread_context_enabled { remaining } else { 0 }; + thread = resume(&test, &thread).await?; + assert_eq!( + serde_json::to_value( + thread + .conversation_history_snapshot() + .await + .retained_context() + )?, + expected + ); + thread.submit(Op::ThreadRollback { num_turns: 1 }).await?; + wait_for_event(&thread, |event| { + matches!(event, EventMsg::ThreadRolledBack(_)) + }) + .await; + expected["user_messages"] + .as_array_mut() + .expect("expected retained user messages") + .truncate(remaining); + expected["verified_answers"] + .as_array_mut() + .expect("expected retained verified answers") + .truncate(remaining); + assert_eq!( + serde_json::to_value( + compact_and_assert_answers(&test, &thread, &answers[..remaining]).await? + )?, + expected + ); + } + thread = resume(&test, &thread).await?; + assert_eq!( + serde_json::to_value( + thread + .conversation_history_snapshot() + .await + .retained_context() + )?, + expected + ); + thread.shutdown_and_wait().await?; + let requests = response_mock.requests(); + assert_eq!(requests.len(), 4 + rollback_counts.len()); + assert!(requests[3].has_message_with_input_texts("user", |texts| texts == [STEER])); + if history_mode == ThreadHistoryMode::Legacy { + assert!(requests[4].has_message_with_input_texts("user", |texts| texts == [INITIAL])); + assert!(!requests[4].has_message_with_input_texts("user", |texts| texts == [STEER])); + assert!(!requests[5].has_message_with_input_texts("user", |texts| texts == [INITIAL])); + } + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn disabled_capture_stays_incomplete_after_compaction_and_enabled_resume() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let test = test_codex() + .with_config(|config| { + config.experimental_thread_store = ThreadStoreConfig::Local; + config + .features + .disable(Feature::GuardianThreadContext) + .expect("disable instruction capture"); + config + .features + .disable(Feature::TokenBudget) + .expect("use local compaction"); + config.model_provider.name = "Local compaction test provider".to_owned(); + }) + .build_with_auto_env(&server) + .await?; + mount_sse_sequence( + &server, + vec![ + sse(vec![ev_completed("initial-turn")]), + sse(vec![ + ev_assistant_message("summary", "Compacted context."), + ev_completed("compact"), + ]), + sse(vec![ + ev_assistant_message("summary-again", "Compacted context."), + ev_completed("compact-again"), + ]), + ], + ) + .await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Never publish publicly.".to_owned(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let checkpoint = compact_and_assert_answers(&test, &test.codex, &[]).await?; + assert!(!checkpoint.user_messages_complete()); + assert_eq!(checkpoint.ordered_entries().count(), 0); + + let thread_id = test.codex.session_configured().thread_id; + test.codex.shutdown_and_wait().await?; + test.thread_manager.remove_thread(&thread_id).await; + let items: Vec = serde_json::from_value(serde_json::to_value( + load_context(&test, &test.codex).await?, + )?)?; + let mut config = test.config.clone(); + config + .features + .enable(Feature::GuardianThreadContext) + .expect("enable instruction capture on resume"); + let resumed = test + .thread_manager + .resume_thread_with_history( + config, + InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history: Arc::new(items), + rollout_path: None, + }), + test.thread_manager.auth_manager(), + /*parent_trace*/ None, + ClientMcpExtensions::default(), + ) + .await? + .thread; + assert!( + !resumed + .conversation_history_snapshot() + .await + .retained_context() + .expect("resumed retained context") + .user_messages_complete() + ); + assert!( + !compact_and_assert_answers(&test, &resumed, &[]) + .await? + .user_messages_complete() + ); + resumed.shutdown_and_wait().await?; Ok(()) } @@ -183,7 +537,11 @@ 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] { + for feature in [ + Feature::TokenBudget, + Feature::DefaultModeRequestUserInput, + Feature::GuardianThreadContext, + ] { config .features .enable(feature) @@ -341,6 +699,7 @@ async fn retained_answers_cross_real_session_boundaries( Feature::DefaultModeRequestUserInput, Feature::Collab, Feature::MultiAgentV2, + Feature::GuardianThreadContext, ] { config .features diff --git a/codex-rs/guardian-context/src/verified_answers.rs b/codex-rs/guardian-context/src/verified_answers.rs index ef9fcf8d20..4357ea2f4d 100644 --- a/codex-rs/guardian-context/src/verified_answers.rs +++ b/codex-rs/guardian-context/src/verified_answers.rs @@ -17,7 +17,7 @@ pub struct RenderedVerifiedAnswers { /// Renders each response with its original source roles and the existing approximate per-answer budget. pub fn render_verified_answers(context: &RetainedContext) -> RenderedVerifiedAnswers { - let mut complete = context.is_complete(); + let mut complete = context.verified_answers_complete(); let mut fragments = Vec::new(); for answer in context.verified_answers() { if let Some(text) = render_verified_answer(answer) { diff --git a/codex-rs/history/src/lib.rs b/codex-rs/history/src/lib.rs index 734ec8e673..a77b1a2e66 100644 --- a/codex-rs/history/src/lib.rs +++ b/codex-rs/history/src/lib.rs @@ -161,7 +161,9 @@ mod guardian_history; mod retained_context; pub use retained_context::RetainedContext; +pub use retained_context::RetainedContextEntry; pub use retained_context::RetainedContextEvent; +pub use retained_context::RetainedUserMessage; pub use retained_context::VerifiedAnswer; pub use retained_context::VerifiedQuestionAnswer; mod rollout_payload; diff --git a/codex-rs/history/src/retained_context.rs b/codex-rs/history/src/retained_context.rs index 97ed607b5d..350dbdf84b 100644 --- a/codex-rs/history/src/retained_context.rs +++ b/codex-rs/history/src/retained_context.rs @@ -7,9 +7,9 @@ use schemars::JsonSchema; use serde::Deserialize; use serde::Serialize; -const MAX_ANSWERS: usize = 8; -const MAX_ANSWER_BYTES: usize = 16_384; -const MAX_ANSWERS_BYTES: usize = 65_536; +const MAX_FAMILY_RECORDS: usize = 8; +const MAX_RECORD_BYTES: usize = 16_384; +const MAX_FAMILY_BYTES: usize = 65_536; /// Original assistant question and host-verified user reply, never an inferred permission. #[derive(Clone, PartialEq, Serialize, Deserialize, JsonSchema)] @@ -27,6 +27,39 @@ pub struct VerifiedAnswer { pub questions: Vec, } +/// Original user instruction, retained outside model summarization for delegated review. +/// Non-text input is not reconstructed as text; missing evidence remains explicit. +#[derive(Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct RetainedUserMessage { + pub turn_id: String, + pub message_id: Option, + pub text: String, + pub complete: bool, +} + +impl std::fmt::Debug for RetainedUserMessage { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("RetainedUserMessage") + .field("complete", &self.complete) + .finish_non_exhaustive() + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +struct Ordered { + #[serde(default)] + order: u64, + #[serde(flatten)] + value: T, +} + +/// Borrowed host evidence in original arrival order, across retained families. +pub enum RetainedContextEntry<'a> { + UserMessage(&'a RetainedUserMessage), + VerifiedAnswer(&'a VerifiedAnswer), +} + impl std::fmt::Debug for VerifiedAnswer { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter @@ -47,9 +80,44 @@ pub enum RetainedContextEvent { /// Facts live until their source instruction is rolled back; compaction does not expire them. #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, JsonSchema)] pub struct RetainedContext { - verified_answers: VecDeque, + verified_answers: VecDeque>, /// Lost evidence cannot be treated as a complete authorization history. - incomplete: bool, + #[serde(rename = "incomplete")] + verified_answers_incomplete: bool, + #[serde(default)] + user_messages: VecDeque>, + /// Old checkpoints did not preserve user restrictions for delegated review. + #[serde(default = "legacy_user_messages_incomplete")] + user_messages_incomplete: bool, + #[serde(default)] + next_order: u64, +} + +fn legacy_user_messages_incomplete() -> bool { + true +} + +fn bound_family(items: &mut VecDeque>, incomplete: &mut bool) { + while items.len() > MAX_FAMILY_RECORDS + || serde_json::to_vec(items).map_or(true, |bytes| bytes.len() > MAX_FAMILY_BYTES) + { + items.pop_front(); + *incomplete = true; + } +} + +impl RetainedUserMessage { + fn bound(&mut self) { + if serde_json::to_vec(self).map_or(true, |bytes| bytes.len() > MAX_RECORD_BYTES) { + self.text.clear(); + self.complete = false; + self.turn_id + .truncate(self.turn_id.floor_char_boundary(1_024)); + if let Some(id) = &mut self.message_id { + id.truncate(id.floor_char_boundary(1_024)); + } + } + } } impl RetainedContextEvent { @@ -57,7 +125,7 @@ impl RetainedContextEvent { pub fn bound(&mut self) { match self { Self::VerifiedAnswer(answer) => { - if serde_json::to_vec(answer).map_or(true, |bytes| bytes.len() > MAX_ANSWER_BYTES) { + 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. answer @@ -74,15 +142,67 @@ impl RetainedContextEvent { impl RetainedContext { pub fn verified_answers(&self) -> impl DoubleEndedIterator { - self.verified_answers.iter() + self.verified_answers.iter().map(|entry| &entry.value) } - pub fn is_complete(&self) -> bool { - !self.incomplete + pub fn verified_answers_complete(&self) -> bool { + !self.verified_answers_incomplete && self .verified_answers .iter() - .all(|answer| !answer.questions.is_empty()) + .all(|answer| !answer.value.questions.is_empty()) + } + + pub fn user_messages_complete(&self) -> bool { + !self.user_messages_incomplete + && self + .user_messages + .iter() + .all(|message| message.value.complete) + } + + /// A skipped instruction leaves a gap that later checkpoints must preserve. + pub fn mark_user_messages_incomplete(&mut self) { + self.user_messages_incomplete = true; + } + + pub fn ordered_entries(&self) -> impl DoubleEndedIterator> { + let mut entries = self + .verified_answers + .iter() + .map(|entry| { + ( + entry.order, + RetainedContextEntry::VerifiedAnswer(&entry.value), + ) + }) + .chain( + self.user_messages + .iter() + .map(|entry| (entry.order, RetainedContextEntry::UserMessage(&entry.value))), + ) + .collect::>(); + entries.sort_by_key(|(order, _)| *order); + 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) { + message.bound(); + if let Some(index) = self.user_messages.iter().position(|entry| { + message.message_id.is_some() && entry.value.message_id == message.message_id + }) { + if self.user_messages[index].value == message { + return; + } + self.user_messages.remove(index); + } + self.user_messages.push_back(Ordered { + order: self.next_order, + value: message, + }); + self.next_order = self.next_order.saturating_add(1); + bound_family(&mut self.user_messages, &mut self.user_messages_incomplete); } /// Same-event delivery is idempotent; changed contents replace that source's record. @@ -92,40 +212,81 @@ impl RetainedContext { match event { RetainedContextEvent::VerifiedAnswer(answer) => { if let Some(index) = self.verified_answers.iter().position(|existing| { - existing.turn_id == answer.turn_id && existing.call_id == answer.call_id + existing.value.turn_id == answer.turn_id + && existing.value.call_id == answer.call_id }) { - if self.verified_answers[index] == answer { + if self.verified_answers[index].value == answer { return false; } self.verified_answers.remove(index); } - self.verified_answers.push_back(answer); - while self.verified_answers.len() > MAX_ANSWERS - || serde_json::to_vec(&self.verified_answers) - .map_or(true, |bytes| bytes.len() > MAX_ANSWERS_BYTES) - { - self.verified_answers.pop_front(); - self.incomplete = true; - } + self.verified_answers.push_back(Ordered { + order: self.next_order, + value: answer, + }); + self.next_order = self.next_order.saturating_add(1); + bound_family( + &mut self.verified_answers, + &mut self.verified_answers_incomplete, + ); } } true } /// Restoring a saved thread must not bypass the live storage limits. - pub fn restore(&mut self, checkpoint: &Self) { - *self = Self { - incomplete: checkpoint.incomplete, + /// A missing checkpoint cannot establish complete historical user instructions. + pub fn restore(&mut self, checkpoint: Option<&Self>) { + *self = checkpoint.cloned().unwrap_or_else(|| Self { + user_messages_incomplete: true, ..Self::default() - }; - for answer in &checkpoint.verified_answers { - self.record(&RetainedContextEvent::VerifiedAnswer(answer.clone())); + }); + for entry in &mut self.verified_answers { + let mut event = RetainedContextEvent::VerifiedAnswer(entry.value.clone()); + event.bound(); + let RetainedContextEvent::VerifiedAnswer(answer) = event; + entry.value = answer; + self.next_order = self.next_order.max(entry.order.saturating_add(1)); } + for entry in &mut self.user_messages { + entry.value.bound(); + self.next_order = self.next_order.max(entry.order.saturating_add(1)); + } + bound_family( + &mut self.verified_answers, + &mut self.verified_answers_incomplete, + ); + bound_family(&mut self.user_messages, &mut self.user_messages_incomplete); } - /// Retains facts whose sources survive a lifecycle change. Compaction alone is not removal. - pub fn retain_answers(&mut self, keep: impl FnMut(&VerifiedAnswer) -> bool) { - self.verified_answers.retain(keep); + /// Keeps legacy answers whose source calls survive when no retained instruction boundary exists. + pub fn retain_answers(&mut self, mut keep: impl FnMut(&VerifiedAnswer) -> bool) { + self.verified_answers.retain(|answer| keep(&answer.value)); + } + + /// Rolls back at the original user-message boundary, including facts delivered after it. + /// 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) + }) { + self.verified_answers.retain(|entry| entry.order < order); + self.user_messages.retain(|entry| entry.order < order); + return; + } + self.user_messages_incomplete |= first_removed_message_id.is_some() + || self + .user_messages + .iter() + .any(|message| turn_ids.contains(&message.value.turn_id.as_str())); + self.verified_answers + .retain(|answer| !turn_ids.contains(&answer.value.turn_id.as_str())); + self.user_messages + .retain(|message| !turn_ids.contains(&message.value.turn_id.as_str())); } } diff --git a/codex-rs/history/src/retained_context_tests.rs b/codex-rs/history/src/retained_context_tests.rs index 5129daae0a..0e945c47e2 100644 --- a/codex-rs/history/src/retained_context_tests.rs +++ b/codex-rs/history/src/retained_context_tests.rs @@ -1,18 +1,29 @@ use super::*; use pretty_assertions::assert_eq; -#[test] -fn answer_lifecycle_is_bounded_idempotent_and_checkpointed() { - let mut context = RetainedContext::default(); - let first = RetainedContextEvent::VerifiedAnswer(VerifiedAnswer { +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(), }], - }); + }) +} + +#[test] +fn retained_evidence_preserves_order_through_checkpoint_and_rollback() { + let mut context = RetainedContext::default(); + 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, + }); let snapshot = context.clone(); assert!(!context.record(&first)); assert_eq!(context, snapshot); @@ -20,8 +31,34 @@ fn answer_lifecycle_is_bounded_idempotent_and_checkpointed() { serde_json::from_str(&serde_json::to_string(&context).expect("retained answer fixture")) .expect("retained answer fixture"); let mut restored = RetainedContext::default(); - restored.restore(&checkpoint); + restored.restore(Some(&checkpoint)); assert_eq!(restored, snapshot); + assert_eq!( + restored + .ordered_entries() + .map(|entry| match entry { + RetainedContextEntry::VerifiedAnswer(answer) => answer.questions[0].answer.as_str(), + RetainedContextEntry::UserMessage(message) => message.text.as_str(), + }) + .collect::>(), + vec!["Yes, but never publicly.", "Do not publish after all."] + ); + restored.rollback(&["revocation-turn"], Some("revocation")); + assert_eq!( + restored, + RetainedContext { + next_order: snapshot.next_order, + ..before_restriction + } + ); +} + +#[test] +fn retained_families_enforce_storage_limits_without_changing_snapshots() { + let mut context = RetainedContext::default(); + let first = publish_answer(); + context.record(&first); + let snapshot = context.clone(); for index in 2..=10 { context.record(&RetainedContextEvent::VerifiedAnswer(VerifiedAnswer { @@ -33,17 +70,18 @@ fn answer_lifecycle_is_bounded_idempotent_and_checkpointed() { }], })); } - assert!(!context.is_complete()); - assert_eq!(context.verified_answers().count(), MAX_ANSWERS); - context.retain_answers(|answer| answer.turn_id != "turn-2"); + assert!(!context.verified_answers_complete()); + assert_eq!(context.verified_answers().count(), MAX_FAMILY_RECORDS); + context.rollback(&["turn-2"], /*first_removed_message_id*/ None); assert_eq!(context.verified_answers().count(), 0); - assert!(!context.is_complete()); + assert!(!context.verified_answers_complete()); + let mut restored = snapshot.clone(); let mut oversized = first.clone(); let RetainedContextEvent::VerifiedAnswer(answer) = &mut oversized; - answer.questions[0].answer = "a".repeat(MAX_ANSWER_BYTES); + answer.questions[0].answer = "a".repeat(MAX_RECORD_BYTES); restored.record(&oversized); - assert!(!restored.is_complete()); + assert!(!restored.verified_answers_complete()); assert!( restored .verified_answers() @@ -61,4 +99,54 @@ fn answer_lifecycle_is_bounded_idempotent_and_checkpointed() { .answer, "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, + }); + } + assert!(!restored.user_messages_complete()); + assert_eq!( + restored + .ordered_entries() + .filter(|entry| matches!(entry, RetainedContextEntry::UserMessage(_))) + .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 Some(RetainedContextEntry::UserMessage(message)) = restored.ordered_entries().next_back() + else { + panic!("latest user evidence"); + }; + assert_eq!((&message.text, message.complete), (&String::new(), false)); +} + +#[test] +fn legacy_checkpoints_mark_user_messages_incomplete() { + let RetainedContextEvent::VerifiedAnswer(answer) = publish_answer(); + let mut wire = serde_json::json!({ + "verified_answers": [answer], "incomplete": false + }); + let legacy: RetainedContext = + serde_json::from_value(wire.clone()).expect("legacy retained-answer checkpoint"); + assert!(legacy.verified_answers_complete()); + assert!(!legacy.user_messages_complete()); + + // Retain the original wire key even though the internal field names its family. + wire["verified_answers"][0]["order"] = serde_json::json!(0); + wire["user_messages"] = serde_json::json!([]); + wire["user_messages_incomplete"] = serde_json::json!(true); + wire["next_order"] = serde_json::json!(0); + assert_eq!(serde_json::to_value(&legacy).unwrap(), wire); + + let mut restored = RetainedContext::default(); + restored.restore(/*checkpoint*/ None); + assert!(!restored.user_messages_complete()); } diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 71d67cc63a..390b1d08b6 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -29,6 +29,7 @@ pub use codex_history::CompactedItem; pub use codex_history::InitialHistory; pub use codex_history::ResponseItemEnvelope; pub use codex_history::ResumedHistory; +pub use codex_history::RetainedContextEntry; pub use codex_history::RetainedContextEvent; pub use codex_history::RolloutItem; pub use codex_history::RolloutLine; 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 44bd5c6937..b16539ea7e 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 @@ -4,16 +4,19 @@ //! Most records happen to be ordered that way, but late completion events can target an older //! surviving turn after a newer turn has started. This planner keeps compact per-record ownership //! metadata for SQLite visibility, then combines it with `rollback_replay`'s cold-resume answer -//! before the writer makes its second streaming pass. -//! Retained answers follow their source calls, not later steers sharing the same turn ID. +//! before the writer makes its second streaming pass. Retained checkpoints are reduced when +//! rollback is encountered, never by reapplying old rollbacks to later checkpoints. use std::collections::HashMap; +use std::collections::HashSet; +use codex_protocol::ResponseItemId; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::UserMessageEvent; use codex_rollout::CompactedItem; +use codex_rollout::RetainedContextEntry; use codex_rollout::RolloutItem; use codex_rollout::RolloutLine; @@ -36,6 +39,12 @@ struct PendingUserResponse { content: Vec, } +struct InstructionBoundary { + record_index: usize, + message_id: Option, + alive: bool, +} + /// Compact plan keyed by parsed source-record index. pub(super) struct RollbackPlan { record_boundaries: Vec>, @@ -79,7 +88,7 @@ impl RollbackPlan { /// Streaming builder for RollbackPlan. pub(super) struct RollbackPlanner { record_boundaries: Vec>, - boundary_alive: Vec, + boundaries: Vec, boundary_stack: Vec, active_turn_id: Option, pending_turn_records: Vec, @@ -88,6 +97,7 @@ pub(super) struct RollbackPlanner { pending_delivery_boundary: Option, turn_boundaries: HashMap, call_boundaries: HashMap<(String, String), Option>, + retained_fact_sources: Vec<(usize, String)>, compactions: Vec, model_replay: ModelReplayPlanner, } @@ -96,7 +106,7 @@ impl RollbackPlanner { pub(super) fn new() -> Self { Self { record_boundaries: Vec::new(), - boundary_alive: Vec::new(), + boundaries: Vec::new(), boundary_stack: Vec::new(), active_turn_id: None, pending_turn_records: Vec::new(), @@ -105,6 +115,7 @@ impl RollbackPlanner { pending_delivery_boundary: None, turn_boundaries: HashMap::new(), call_boundaries: HashMap::new(), + retained_fact_sources: Vec::new(), compactions: Vec::new(), model_replay: ModelReplayPlanner::new(), } @@ -139,8 +150,10 @@ impl RollbackPlanner { RolloutItem::ResponseItem(response) => { if let Some(boundary) = paired_delivery_boundary { self.record_boundaries[index] = Some(boundary); + self.boundaries[boundary].message_id = response.id().cloned(); } else if rollback::counts_as_boundary(&response.item) { let boundary = self.start_boundary(index); + self.boundaries[boundary].message_id = response.id().cloned(); if let ResponseItem::Message { role, content, .. } = &response.item && role == "user" { @@ -247,6 +260,8 @@ impl RollbackPlanner { } self.call_boundaries .insert(source, self.record_boundaries[index]); + self.retained_fact_sources + .push((index, answer.turn_id.clone())); } RolloutItem::SecurityRiskScore(_) => self.record_boundaries[index] = None, } @@ -257,34 +272,19 @@ impl RollbackPlanner { pub(super) fn finish(self) -> RollbackPlan { let RollbackPlanner { record_boundaries, - boundary_alive, + boundaries, compactions, model_replay, - turn_boundaries, - call_boundaries, .. } = self; let replay_anchor = model_replay.finish().empty_replacement_history_compaction; - let removed_turns = turn_boundaries - .iter() - .filter(|(_, boundary)| !boundary_alive[**boundary]) - .map(|(turn_id, _)| turn_id.as_str()) + let boundary_alive = boundaries + .into_iter() + .map(|boundary| boundary.alive) .collect::>(); let compacted_items = compactions .into_iter() .filter_map(|mut frame| { - // Migration removes rollback markers. Checkpoints must therefore carry - // the same surviving facts as their standalone source events. - if let Some(context) = &mut frame.item.retained_context { - context.retain_answers(|answer| { - call_boundaries - .get(&(answer.turn_id.clone(), answer.call_id.clone())) - .map_or_else( - || !removed_turns.contains(&answer.turn_id.as_str()), - |boundary| boundary.is_none_or(|boundary| boundary_alive[boundary]), - ) - }); - } if Some(frame.record_index) == replay_anchor { frame.item.replacement_history = Some(Vec::new()); frame.item.mcp_resource_origins = None; @@ -304,8 +304,12 @@ impl RollbackPlanner { } fn start_boundary(&mut self, index: usize) -> usize { - let boundary = self.boundary_alive.len(); - self.boundary_alive.push(true); + let boundary = self.boundaries.len(); + self.boundaries.push(InstructionBoundary { + record_index: index, + message_id: None, + alive: true, + }); let had_prior_boundary = !self.boundary_stack.is_empty(); if had_prior_boundary { for pending_index in self.pending_context_records.drain(..) { @@ -348,16 +352,72 @@ impl RollbackPlanner { return Ok(()); } let depth_before = self.boundary_stack.len(); + let mut removed_boundaries = HashSet::new(); + let mut first_removed_boundary = None; for _ in 0..count { let Some(boundary) = self.boundary_stack.pop() else { break; }; - self.boundary_alive[boundary] = false; + self.boundaries[boundary].alive = false; + removed_boundaries.insert(boundary); + first_removed_boundary = Some(boundary); + } + if let Some(boundary) = first_removed_boundary { + let source = &self.boundaries[boundary]; + let removed_turns = self + .turn_boundaries + .iter() + .filter(|(_, boundary)| removed_boundaries.contains(*boundary)) + .map(|(turn_id, _)| turn_id.as_str()) + .chain( + self.retained_fact_sources + .iter() + .filter(|(index, _)| { + self.record_boundaries[*index] + .is_some_and(|boundary| removed_boundaries.contains(&boundary)) + }) + .map(|(_, turn_id)| turn_id.as_str()), + ) + .collect::>(); + // A checkpoint before this instruction cannot contain evidence from it. + // 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) + { + if let Some(context) = &mut frame.item.retained_context { + if context + .ordered_entries() + .any(|entry| matches!(entry, RetainedContextEntry::UserMessage(_))) + { + context.rollback( + &removed_turns, + source.message_id.as_ref().map(ResponseItemId::as_str), + ); + } else { + // Checkpoints written without instruction retention keep the legacy + // source-call boundary, including answers before a same-turn steer. + context.retain_answers(|answer| { + self.call_boundaries + .get(&(answer.turn_id.clone(), answer.call_id.clone())) + .map_or_else( + || !removed_turns.contains(&answer.turn_id.as_str()), + |boundary| { + boundary + .is_none_or(|boundary| self.boundaries[boundary].alive) + }, + ) + }); + } + } + } } let compaction_index = self.compactions.iter().rposition(|frame| { frame .owner - .is_none_or(|boundary| self.boundary_alive[boundary]) + .is_none_or(|boundary| self.boundaries[boundary].alive) }); if let Some(compaction_index) = compaction_index { let frame = &mut self.compactions[compaction_index]; 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 6ca78276ea..0758db8163 100644 --- a/codex-rs/thread-store/src/local/rollout_migration_tests.rs +++ b/codex-rs/thread-store/src/local/rollout_migration_tests.rs @@ -11,6 +11,7 @@ use codex_extension_items::ExtensionItem; use codex_extension_items::image_generation::ImageGenerationFailure; use codex_extension_items::image_generation::ImageGenerationItem; use codex_protocol::AgentPath; +use codex_protocol::ResponseItemId; use codex_protocol::ThreadId; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::items::ReasoningItem; @@ -18,6 +19,8 @@ use codex_protocol::items::TurnItem; use codex_protocol::mcp::McpResourceOrigin; use codex_protocol::mcp::McpResourceOriginCheckpoint; use codex_protocol::models::ContentItem; +use codex_protocol::models::ContentItemKind; +use codex_protocol::models::InternalChatMessageMetadataPassthrough; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::AgentMessageEvent; use codex_protocol::protocol::AskForApproval; @@ -1267,6 +1270,139 @@ async fn migration_preserves_reverse_replay_anchor_after_pre_compaction_rollback assert!(replacement_history.is_empty()); } +#[tokio::test] +async fn migration_preserves_same_turn_evidence_across_rollback_checkpoints() { + const INITIAL: &str = "Never publish publicly."; + const STEER: &str = "Also inspect the README."; + let home = TempDir::new().expect("create Codex home"); + let thread_id = ThreadId::new(); + let [initial, steer] = + [("initial", INITIAL), ("steer", STEER)].map(|(id, text)| ResponseItem::Message { + id: Some(ResponseItemId::with_suffix("msg", id)), + role: "user".to_owned(), + content: vec![ContentItem::InputText { + text: text.to_owned(), + }], + phase: None, + internal_chat_message_metadata_passthrough: Some( + InternalChatMessageMetadataPassthrough { + turn_id: Some("shared-turn".to_owned()), + content_item_kinds: Some(vec![ContentItemKind("user.text".to_owned())]), + ..Default::default() + }, + ), + }); + let RolloutItem::Compacted(mut checkpoint) = compacted(vec![initial.clone(), steer.clone()]) + else { + unreachable!("compacted helper creates a checkpoint"); + }; + 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} + ], + "verified_answers": [ + {"order": 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 + }); + checkpoint.retained_context = + Some(serde_json::from_value(retained.clone()).expect("retained fixture")); + let answers = checkpoint + .retained_context + .as_ref() + .expect("retained checkpoint") + .verified_answers() + .cloned() + .map(codex_rollout::RetainedContextEvent::VerifiedAnswer) + .collect::>(); + let mut expected = retained; + expected["user_messages"] + .as_array_mut() + .expect("user messages") + .truncate(/*len*/ 1); + expected["verified_answers"] + .as_array_mut() + .expect("verified answers") + .truncate(/*len*/ 1); + let mut before_expected = expected.clone(); + before_expected["next_order"] = json!(2); + 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")); + let mut after_rollback = before_steer.clone(); + after_rollback.retained_context = + Some(serde_json::from_value(expected.clone()).expect("post-rollback checkpoint")); + let path = write_rollout( + home.path(), + thread_id, + SessionSource::Cli, + vec![ + started("shared-turn"), + rollout_response_item(initial.clone()), + user_message(INITIAL), + RolloutItem::RetainedContext(answers[0].clone()), + RolloutItem::Compacted(before_steer), + rollout_response_item(steer), + user_message(STEER), + RolloutItem::RetainedContext(answers[1].clone()), + completed("shared-turn"), + started("compaction-turn"), + RolloutItem::Compacted(checkpoint), + completed("compaction-turn"), + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns: 1, + })), + // This checkpoint already contains the correct live rollback result. + started("post-rollback-compaction"), + RolloutItem::Compacted(after_rollback), + completed("post-rollback-compaction"), + ], + ); + let store = indexed_store(home.path()).await; + store + .migrate_rollouts(apply_options()) + .await + .expect("migrate same-turn rollback"); + let migrated = read_rollout(&path); + let checkpoints = migrated + .iter() + .filter_map(|line| match &line.item { + RolloutItem::Compacted(item) => Some(item), + _ => None, + }) + .collect::>(); + assert_eq!( + checkpoints + .iter() + .map( + |checkpoint| serde_json::to_value(&checkpoint.retained_context) + .expect("retained context") + ) + .collect::>(), + vec![before_expected, expected.clone(), expected] + ); + assert_eq!( + checkpoints + .last() + .expect("latest checkpoint") + .replacement_history, + Some(vec![initial.into()]) + ); + assert_eq!( + migrated + .into_iter() + .filter_map(|line| match line.item { + RolloutItem::RetainedContext(event) => Some(event), + _ => None, + }) + .collect::>(), + vec![answers[0].clone()] + ); +} + #[tokio::test] async fn migration_keeps_empty_replay_anchor_from_rolled_back_turn() { let home = TempDir::new().expect("create Codex home");