openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-28 18:58:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Retain user instructions in guardian thread context (#42762)

Browse Source 
## Why

Guardian review needs the original user instructions alongside verified answers, even after compaction, resume, or rollback. A summarized or legacy checkpoint cannot establish a complete instruction history.

## What changed

- Capture text user messages in retained context when `guardian_thread_context` is enabled, preserving their message identity through compaction and replay.
- Keep retained instructions and verified answers in arrival order, with independent size and completeness tracking for each family.
- Roll back retained evidence at the exact user-message boundary, including steers that share a turn ID, and apply the same semantics during rollout migration.
- Treat checkpoints without retained user messages as incomplete while preserving their verified answers.

## Testing

- Cover capture, compaction, resume, rollback, child forks, legacy checkpoints, storage bounds, and rollout migration.

GitOrigin-RevId: 72706731b54e5ac8645db61b446d2401afb42235
This commit is contained in:
felixxia-oai
2026-09-04 13:24:42 +00:00
committed by copyberry
parent cc4b8bdeb8
commit 80d7ca34bc
18 changed files with 1058 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst
Generated
BIN
View File
Binary file not shown.
codex-rs/core/src/agent/control_tests.rs
+3 -4
View File
@@ -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"
codex-rs/core/src/compact.rs
+26 -2
View File
@@ -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<String> {
#[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<ResponseItemId>,
message: String,
internal_chat_message_metadata_passthrough: Option<InternalChatMessageMetadataPassthrough>,
harness_metadata: Option<CodexHarnessMetadata>,
@@ -540,12 +550,24 @@ pub(crate) fn collect_user_messages(items: &[ResponseItem]) -> Vec<CompactedUser
.collect()
}
pub(crate) enum CompactedMessageIdentity {
Preserve,
Regenerate,
}
pub(crate) fn collect_annotated_user_messages(
items: &[ResponseItemEnvelope],
identity: CompactedMessageIdentity,
) -> Vec<CompactedUserMessage> {
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(),
codex-rs/core/src/compact_tests.rs
+14 -3
View File
@@ -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(),
codex-rs/core/src/context_manager/history.rs
+64 -13
View File
@@ -61,6 +61,8 @@ pub(crate) struct ContextManager {
review_history: Option<TranscriptHistory>,
/// Host facts independent of the model window; snapshots share immutable state.
retained_context: Arc<RetainedContext>,
/// 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::<Vec<_>>()
.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::<Vec<_>>();
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;
codex-rs/core/src/context_manager/history_tests.rs
+49
View File
@@ -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]
codex-rs/core/src/session/mod.rs
+3 -1
View File
@@ -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);
}
codex-rs/core/src/session/rollout_reconstruction.rs
+14 -6
View File
@@ -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)) => {
codex-rs/core/src/session/rollout_reconstruction_tests.rs
+1
View File
@@ -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()),
codex-rs/core/src/session/session.rs
+3
View File
@@ -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
codex-rs/core/tests/suite/guardian_retained_context.rs
+364 -5
View File
@@ -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<RetainedContext> {
// 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::<Vec<_>>(),
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::<Vec<_>>();
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::<Vec<_>>();
let ordered_answers = answers
.iter()
.enumerate()
.map(|(index, answer)| {
let mut value = json!(answer);
value["order"] = json!(index * 2 + 1);
value
})
.collect::<Vec<_>>();
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<RolloutItem> = 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
codex-rs/guardian-context/src/verified_answers.rs
+1 -1
View File
@@ -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) {
codex-rs/history/src/lib.rs
+2
View File
@@ -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;
codex-rs/history/src/retained_context.rs
+190 -29
View File
@@ -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<VerifiedQuestionAnswer>,
}
/// 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<String>,
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<T> {
#[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<VerifiedAnswer>,
verified_answers: VecDeque<Ordered<VerifiedAnswer>>,
/// Lost evidence cannot be treated as a complete authorization history.
incomplete: bool,
#[serde(rename = "incomplete")]
verified_answers_incomplete: bool,
#[serde(default)]
user_messages: VecDeque<Ordered<RetainedUserMessage>>,
/// 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<T: Serialize>(items: &mut VecDeque<Ordered<T>>, 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<Item = &VerifiedAnswer> {
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<Item = RetainedContextEntry<'_>> {
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::<Vec<_>>();
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()));
}
}
codex-rs/history/src/retained_context_tests.rs
+100 -12
View File
@@ -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<_>>(),
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());
}
codex-rs/rollout/src/lib.rs
+1
View File
@@ -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;
codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs
+87 -27
View File
@@ -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<ContentItem>,
}
struct InstructionBoundary {
record_index: usize,
message_id: Option<ResponseItemId>,
alive: bool,
}
/// Compact plan keyed by parsed source-record index.
pub(super) struct RollbackPlan {
record_boundaries: Vec<Option<usize>>,
@@ -79,7 +88,7 @@ impl RollbackPlan {
/// Streaming builder for RollbackPlan.
pub(super) struct RollbackPlanner {
record_boundaries: Vec<Option<usize>>,
boundary_alive: Vec<bool>,
boundaries: Vec<InstructionBoundary>,
boundary_stack: Vec<usize>,
active_turn_id: Option<String>,
pending_turn_records: Vec<usize>,
@@ -88,6 +97,7 @@ pub(super) struct RollbackPlanner {
pending_delivery_boundary: Option<usize>,
turn_boundaries: HashMap<String, usize>,
call_boundaries: HashMap<(String, String), Option<usize>>,
retained_fact_sources: Vec<(usize, String)>,
compactions: Vec<CompactionFrame>,
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::<Vec<_>>();
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::<Vec<_>>();
// 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];
codex-rs/thread-store/src/local/rollout_migration_tests.rs
+136
View File
@@ -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::<Vec<_>>();
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::<Vec<_>>();
assert_eq!(
checkpoints
.iter()
.map(
|checkpoint| serde_json::to_value(&checkpoint.retained_context)
.expect("retained context")
)
.collect::<Vec<_>>(),
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<_>>(),
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");