Preserve retained assistant revisions across rollout replay (#60839)

## Why

Replaying delivered assistant messages allocates new retained-source revisions instead of preserving those captured in live history.

## What changed

- Capture retained sources when recording assistant deliveries, persist them in rollout metadata, and reuse their revisions during replay.
- Add optional `next_item_ordinal` compaction metadata and a `next_model_id_ordinals` helper to recover item and window allocation positions from checkpoints and subsequent records. Ignore allocation metadata that refers to other threads.

## Testing

Extend tests to compare live and replayed delivery sources, preserve revisions through an older reader's rewrite, and verify ordinal recovery from checkpoints and subsequent local items.

GitOrigin-RevId: cfdeebe8e13b1700b34b9c98696880a90a99ace5
This commit is contained in:
Ruslan Nigmatullin
2026-11-11 18:56:47 +00:00
committed by copyberry
parent e58bf8f523
commit 9e6d1c5278
21 changed files with 238 additions and 25 deletions
+1
View File
@@ -3850,6 +3850,7 @@ dependencies = [
"codex-protocol",
"codex-protocol-test-support",
"codex-utils-absolute-path",
"codex-utils-model-facing-id",
"pretty_assertions",
"schemars 0.8.22",
"serde",
+3
View File
@@ -2716,6 +2716,7 @@ async fn spawn_agent_fork_sanitizes_inherited_compaction_metadata() {
let (parent_thread_id, parent_thread) = harness.start_paginated_thread().await;
let parent_resume_metadata = codex_history::CompactionResumeMetadata {
next_item_ordinal: None,
rejected_audio_file_ids: Default::default(),
dynamic_tools: None,
multi_agent_version: Some(MultiAgentVersion::V2),
@@ -3881,6 +3882,7 @@ async fn spawn_agent_full_fork_legacy_compaction_rebuilds_child_instructions_onc
compaction_response_id: None,
latest_token_usage_record: None,
resume_metadata: Some(codex_history::CompactionResumeMetadata {
next_item_ordinal: None,
rejected_audio_file_ids: Default::default(),
dynamic_tools: None,
multi_agent_version: Some(MultiAgentVersion::V2),
@@ -4017,6 +4019,7 @@ async fn spawn_agent_full_fork_legacy_compaction_rebuilds_child_instructions_onc
assert_eq!(
inherited_state,
Some(codex_history::CompactionResumeMetadata {
next_item_ordinal: None,
rejected_audio_file_ids: Default::default(),
dynamic_tools: None,
multi_agent_version: Some(MultiAgentVersion::V2),
+10 -1
View File
@@ -389,7 +389,16 @@ impl ContextManager {
}
pub(crate) fn record_retained_context(&mut self, event: &RetainedContextEvent) -> bool {
if !Arc::make_mut(&mut self.retained_context).record(event, &self.item_id_generator) {
self.record_retained_context_for_persistence(&mut event.clone())
}
pub(crate) fn record_retained_context_for_persistence(
&mut self,
event: &mut RetainedContextEvent,
) -> bool {
if !Arc::make_mut(&mut self.retained_context)
.record_and_capture(event, &self.item_id_generator)
{
return false;
}
match event {
@@ -203,6 +203,7 @@ fn delivered_assistant_context_invalidates_reviews_without_changing_authorizatio
let before_delivery = history.conversation_history_snapshot();
let authorization = before_delivery.user_message_revision();
let event = RetainedContextEvent::DeliveredAssistantMessage {
retained_source: None,
message: codex_history::RetainedUserMessage {
origin: codex_history::UserInputOrigin::User,
turn_id: "turn".to_owned(),
+1
View File
@@ -4959,6 +4959,7 @@ impl Session {
compaction_response_id: metadata.compaction_response_id,
latest_token_usage_record: state.latest_token_usage_record.clone(),
resume_metadata: Some(CompactionResumeMetadata {
next_item_ordinal: None,
rejected_audio_file_ids: state.history.rejected_audio_file_ids(),
dynamic_tools: Some(Box::new(DynamicToolsSnapshot {
thread_id: self.thread_id(),
@@ -156,7 +156,6 @@ impl Session {
}
pub(crate) async fn record_retained_context(&self, mut event: RetainedContextEvent) {
event.bound();
// Share the checkpoint persistence lock so a fact cannot land on the wrong side
// of the checkpoint/suffix boundary. Ephemeral threads use the same live state.
let _guard = thread_settings::acquire_persistence_lock(self).await;
@@ -165,7 +164,7 @@ impl Session {
.lock()
.await
.history
.record_retained_context(&event)
.record_retained_context_for_persistence(&mut event)
{
self.persist_rollout_items(&[RolloutItem::RetainedContext(event)])
.await;
@@ -211,10 +210,12 @@ impl Session {
let mut state = session.state.lock().await;
let mut event = RetainedContextEvent::DeliveredAssistantMessage {
message,
retained_source: None,
acceptance_order: state.history.reserve_input_order(),
};
event.bound();
state.history.record_retained_context(&event);
state
.history
.record_retained_context_for_persistence(&mut event);
let mut pending = session
.code_mode_message_tasks
.pending_persistence
@@ -1642,6 +1642,7 @@ async fn bounded_replay_matches_full_replay_after_empty_turn_compactions(
compaction_response_id: None,
latest_token_usage_record: None,
resume_metadata: current.then(|| codex_history::CompactionResumeMetadata {
next_item_ordinal: None,
rejected_audio_file_ids: Default::default(),
dynamic_tools: None,
multi_agent_version: None,
+20
View File
@@ -2631,6 +2631,25 @@ async fn inter_agent_communication_waits_for_confirmed_delivery_persistence() {
]
);
let live = session.clone_history().await;
let replayed = session
.reconstruct_history_from_rollout(&replay_turn, &resumed.history)
.await;
let delivery_sources = |context: &codex_history::RetainedContext| {
context
.ordered_entries()
.filter_map(|(_, entry)| match entry {
codex_history::RetainedContextEntry::AssistantMessage(_) => Some(
context
.source(entry)
.expect("confirmed delivery has a revision"),
),
_ => None,
})
.collect::<Vec<_>>()
};
let expected = delivery_sources(live.retained_context());
assert_eq!(expected.len(), 2);
assert_eq!(delivery_sources(&replayed.retained_context), expected);
let communication = live
.annotated_items()
.iter()
@@ -6219,6 +6238,7 @@ async fn compaction_persists_resume_metadata_and_companion_records() {
session.multi_agent_version = std::sync::OnceLock::from(MultiAgentVersion::V2);
let expected = CompactionResumeMetadata {
next_item_ordinal: None,
rejected_audio_file_ids: Default::default(),
dynamic_tools: Some(Box::new(DynamicToolsSnapshot {
thread_id: session.thread_id(),
@@ -130,6 +130,7 @@ Approved action: {"action":{"command":"write /workspace/report.txt"},"outcome":"
acceptance_order: Some(2),
});
let delivered = RolloutItem::RetainedContext(RetainedContextEvent::DeliveredAssistantMessage {
retained_source: None,
message: RetainedUserMessage {
turn_id: "turn-1".into(),
message_id: Some("delivered-1".into()),
+1
View File
@@ -15,6 +15,7 @@ workspace = true
[dependencies]
codex-protocol = { workspace = true }
codex-utils-absolute-path = { workspace = true }
codex-utils-model-facing-id = { workspace = true }
schemars = { workspace = true }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
@@ -18,6 +18,9 @@ use std::sync::Arc;
/// compactions that did not.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
pub struct CompactionResumeMetadata {
/// Next ordinal in the producing thread's seeded item ID sequence at this checkpoint.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub next_item_ordinal: Option<u64>,
/// Immutable file references rejected before this checkpoint.
#[serde(default, skip_serializing_if = "BTreeSet::is_empty")]
pub rejected_audio_file_ids: Arc<BTreeSet<String>>,
+2
View File
@@ -290,6 +290,8 @@ impl JsonSchema for RolloutItem {
}
mod guardian_history;
mod model_id_ordinals;
pub use model_id_ordinals::next_model_id_ordinals;
mod reconciled_retained_context;
mod retained_context;
mod sender_user_messages;
+70
View File
@@ -0,0 +1,70 @@
//! Recovers seeded allocation positions from a compaction checkpoint and subsequent thread items.
use codex_protocol::ResponseItemId;
use codex_utils_model_facing_id::ModelFacingId;
use crate::RetainedContextEvent;
use crate::RolloutItem;
/// Next item and window ordinals. Legacy strings have no allocation provenance.
pub fn next_model_id_ordinals(records: &[RolloutItem]) -> (u64, u64) {
let mut item = None;
let mut window = None;
for record in records {
match record {
RolloutItem::SessionMeta(meta) => {
if let Some(initial) = &meta.meta.context_window {
window = window.max(initial.window_id.allocation().map(ModelFacingId::ordinal));
}
}
RolloutItem::Compacted(checkpoint) => {
if let Some(next) = checkpoint
.resume_metadata
.as_ref()
.and_then(|m| m.next_item_ordinal)
{
item = next.checked_sub(1);
}
if let Some(id) = &checkpoint.window_id {
window = window.max(id.allocation().map(ModelFacingId::ordinal));
}
}
RolloutItem::ResponseItem(envelope) => {
item = item.max(envelope.item.id().and_then(item_ordinal));
// The retained revision is local. Guardian delivery metadata quotes other threads.
if let Some(source) = envelope
.metadata
.as_ref()
.and_then(|m| m.retained_source.as_ref())
{
item = item.max(item_ordinal(&source.revision));
}
}
RolloutItem::InterAgentCommunication(message) => {
item = item.max(message.id.as_ref().and_then(item_ordinal));
}
RolloutItem::RetainedContext(RetainedContextEvent::DeliveredAssistantMessage {
retained_source: Some(source),
..
}) => {
item = item.max(item_ordinal(&source.revision));
}
RolloutItem::DynamicToolsUpdated { .. }
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::TurnContext(_)
| RolloutItem::TokenUsageRecord(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::RetainedContext(_)
| RolloutItem::EventMsg(_)
| RolloutItem::RealtimeItem(_) => {}
}
}
let next = |last: Option<u64>| last.map_or(0, |ordinal| ordinal.wrapping_add(1));
(next(item), next(window))
}
fn item_ordinal(id: &ResponseItemId) -> Option<u64> {
id.allocation()
.map(|allocation| allocation.model_id.ordinal())
}
@@ -37,9 +37,20 @@ impl RetainedContext {
/// Returns the captured source for the original envelope, independently of buffer eviction.
pub fn record_assistant_message(
&mut self,
mut message: RetainedUserMessage,
message: RetainedUserMessage,
source: RetainedInputSource,
item_id_generator: &ItemIdGenerator,
) -> Option<RetainedSource> {
self.record_assistant_message_with_revision(message, source, || {
item_id_generator("retained")
})
}
pub(super) fn record_assistant_message_with_revision(
&mut self,
mut message: RetainedUserMessage,
source: RetainedInputSource,
revision: impl FnOnce() -> codex_protocol::ResponseItemId,
) -> Option<RetainedSource> {
if source == RetainedInputSource::Local(None) {
// Leave unsequenced sources to legacy transcript selection. Missing
@@ -72,10 +83,7 @@ impl RetainedContext {
if empty {
return None;
}
let revision = message
.message_id
.as_ref()
.map(|_| item_id_generator("retained"));
let revision = message.message_id.as_ref().map(|_| revision());
let entry = Ordered {
revision,
inherited,
+30 -12
View File
@@ -179,6 +179,9 @@ pub enum RetainedContextEvent {
#[serde(flatten)]
message: RetainedUserMessage,
acceptance_order: u64,
/// The host-captured evidence revision. Older deliveries only persisted the call ID.
#[serde(default, skip_serializing_if = "Option::is_none")]
retained_source: Option<RetainedSource>,
},
}
@@ -477,16 +480,24 @@ impl RetainedContext {
event: &RetainedContextEvent,
item_id_generator: &ItemIdGenerator,
) -> bool {
let mut event = event.clone();
self.record_and_capture(&mut event.clone(), item_id_generator)
}
/// Records a host event and fills in its retained source for rollout persistence.
pub fn record_and_capture(
&mut self,
event: &mut RetainedContextEvent,
item_id_generator: &ItemIdGenerator,
) -> bool {
event.bound();
match event {
RetainedContextEvent::AcceptedFileAudio { evidence } => {
if !evidence.fits_record_budget()
|| self.accepted_file_audio.as_ref() == Some(&evidence)
|| self.accepted_file_audio.as_ref() == Some(evidence)
{
return false;
}
self.accepted_file_audio = Some(evidence);
self.accepted_file_audio = Some(Arc::clone(evidence));
true
}
RetainedContextEvent::VerifiedAnswer {
@@ -497,17 +508,17 @@ impl RetainedContext {
existing.value.turn_id == answer.turn_id
&& existing.value.call_id == answer.call_id
}) {
if self.verified_answers[index].value == answer {
if self.verified_answers[index].value == *answer {
return false;
}
self.verified_answers.remove(index);
}
let order = self.record_order(acceptance_order);
let order = self.record_order(*acceptance_order);
self.verified_answers.push_back(Ordered {
revision: None,
inherited: false,
order,
value: answer,
value: answer.clone(),
});
bound_family(
&mut self.verified_answers,
@@ -518,19 +529,26 @@ impl RetainedContext {
RetainedContextEvent::DeliveredAssistantMessage {
message,
acceptance_order,
retained_source,
} => {
if self.assistant_messages.iter().any(|entry| {
if let Some(existing) = self.assistant_messages.iter().find(|entry| {
message.message_id.is_some()
&& entry.value.message_id == message.message_id
&& entry.value == message
&& entry.value == *message
&& !entry.inherited
}) {
*retained_source = existing.source(RetainedSourceRole::Assistant);
return false;
}
self.record_assistant_message(
message,
RetainedInputSource::Local(Some(acceptance_order)),
item_id_generator,
*retained_source = self.record_assistant_message_with_revision(
message.clone(),
RetainedInputSource::Local(Some(*acceptance_order)),
|| {
retained_source.as_ref().map_or_else(
|| item_id_generator("retained"),
|source| source.revision.clone(),
)
},
);
true
}
@@ -51,6 +51,7 @@ fn publish_answer() -> RetainedContextEvent {
fn delivered_message(id: &str, text: &str, acceptance_order: u64) -> RetainedContextEvent {
RetainedContextEvent::DeliveredAssistantMessage {
retained_source: None,
message: RetainedUserMessage {
origin: crate::UserInputOrigin::User,
turn_id: "turn-1".to_owned(),
+3
View File
@@ -218,6 +218,7 @@ impl<'a> From<&'a RolloutItem> for RolloutItemWire<'a> {
RolloutItem::RetainedContext(RetainedContextEvent::DeliveredAssistantMessage {
message,
acceptance_order,
retained_source,
}) => {
// Earlier clients preserve this as genuine assistant context, including across compaction.
let text = if message.complete {
@@ -256,6 +257,7 @@ impl<'a> From<&'a RolloutItem> for RolloutItemWire<'a> {
id_allocation: None,
},
metadata: Some(Cow::Owned(CodexHarnessMetadata {
retained_source: retained_source.clone(),
delivered_assistant_message: Some(marker),
user_input_order: Some(*acceptance_order),
..Default::default()
@@ -323,6 +325,7 @@ impl From<RolloutItemWire<'_>> for RolloutItem {
phase: None,
},
acceptance_order,
retained_source: metadata.retained_source.clone(),
},
);
}
+60 -1
View File
@@ -1,5 +1,6 @@
use anyhow::Result;
use codex_protocol::ContextWindowId;
use codex_protocol::ResponseItemId;
use codex_protocol::models::AudioReference;
use codex_protocol::models::ConfigurationReasoning;
use codex_protocol::models::PermissionProfile;
@@ -366,7 +367,8 @@ fn delivered_assistant_rollout_survives_an_earlier_read_and_rewrite() -> Result<
}
for complete in [true, false] {
let original = RetainedContextEvent::DeliveredAssistantMessage {
let mut original = RetainedContextEvent::DeliveredAssistantMessage {
retained_source: None,
message: RetainedUserMessage {
origin: crate::UserInputOrigin::User,
turn_id: "root-turn".to_owned(),
@@ -377,6 +379,14 @@ fn delivered_assistant_rollout_survives_an_earlier_read_and_rewrite() -> Result<
},
acceptance_order: 7,
};
let sequence = std::sync::Mutex::new(
codex_utils_model_facing_id::ModelFacingIdGenerator::new(/*seed*/ 42, "item"),
);
let generator: codex_protocol::ItemIdGenerator = Arc::new(move |kind| {
ResponseItemId::generated(kind, sequence.lock().unwrap().fetch_next())
});
let mut live = RetainedContext::default();
assert!(live.record_and_capture(&mut original, &generator));
let wire = serde_json::to_value(RolloutItem::RetainedContext(original.clone()))?;
let earlier = serde_json::from_value::<EarlierRolloutItem>(wire)?;
let EarlierRolloutItem::ResponseItem { payload, metadata } = &earlier;
@@ -432,6 +442,13 @@ fn delivered_assistant_rollout_survives_an_earlier_read_and_rewrite() -> Result<
panic!("new readers need the model-invisible delivery event");
};
assert_eq!(restored, original);
let mut replayed = RetainedContext::default();
assert!(replayed.record(&restored, &generator));
assert_eq!(replayed, live);
assert_eq!(
crate::next_model_id_ordinals(&[RolloutItem::RetainedContext(restored)]),
(1, 0)
);
}
Ok(())
}
@@ -679,6 +696,7 @@ fn compacted_replacement_history_stores_metadata_in_an_aligned_sidecar() -> Resu
#[test]
fn compacted_resume_metadata_presence_round_trips_empty_values() -> Result<()> {
let resume_metadata = CompactionResumeMetadata {
next_item_ordinal: None,
rejected_audio_file_ids: Default::default(),
dynamic_tools: None,
multi_agent_version: None,
@@ -792,6 +810,7 @@ fn compacted_metadata_remains_compatible_with_legacy_response_item_readers() ->
compaction_response_id: None,
latest_token_usage_record: None,
resume_metadata: Some(CompactionResumeMetadata {
next_item_ordinal: None,
rejected_audio_file_ids: Default::default(),
dynamic_tools: None,
multi_agent_version: Some(MultiAgentVersion::V2),
@@ -1158,6 +1177,46 @@ fn compacted_item_serializes_window_number_and_id() -> Result<()> {
Ok(())
}
#[test]
fn model_id_ordinals_resume_at_the_checkpoint_and_advance_for_local_items() -> Result<()> {
let record = |ordinal| {
json!({
"type": "response_item", "payload": {"type": "message", "role": "assistant", "content": [],
"id": "external", "id_allocation": {"salt": "item", "ordinal": ordinal, "value": 3}}
})
};
let source = |ordinal| {
json!({
"id": {"turn_id": "turn", "message_id": "external", "role": "assistant"},
"revision": {"id": "retained_item007",
"id_allocation": {"prefix": "retained", "salt": "item", "ordinal": ordinal, "value": 7}},
"complete": true
})
};
let mut tail = record(24);
tail["metadata"] = json!({
"retained_source": source(25),
"guardian_sources": [source(300)],
"guardian_review_ids": [{"id": "review_item008",
"id_allocation": {"prefix": "review", "salt": "item", "ordinal": 301, "value": 8}}]
});
let records: Vec<RolloutItem> = serde_json::from_value(json!([
record(19),
{"type": "compacted", "payload": {
"message": "", "window_number": 4,
"window_id": {"salt": "window", "ordinal": 4, "value": 812},
"replacement_history": [record(8)["payload"]],
"resume_metadata": {"next_item_ordinal": 24}
}},
tail
]))?;
assert_eq!(crate::next_model_id_ordinals(&[]), (0, 0));
assert_eq!(crate::next_model_id_ordinals(&records[1..2]), (24, 5));
assert_eq!(crate::next_model_id_ordinals(&records[1..]), (26, 5));
assert_eq!(crate::next_model_id_ordinals(&records), (26, 5));
Ok(())
}
#[test]
fn structured_window_ids_round_trip_in_session_metadata_and_compaction() -> Result<()> {
let uuid = "019020ef-10ff-7000-8000-000000000001";
@@ -95,6 +95,14 @@ async fn stops_at_newest_usable_compaction_and_keeps_companions() {
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("thread id");
let world_state = RolloutItem::WorldState(WorldStateItem::full(Default::default()));
let turn_context = turn_context(home.path(), "companion");
let checkpoint: RolloutItem = serde_json::from_value(serde_json::json!({
"type": "compacted", "payload": {
"message": "latest compaction", "replacement_history": [], "window_number": 2,
"window_id": {"salt": "window", "ordinal": 2, "value": 6},
"resume_metadata": {"next_item_ordinal": 24}
}
}))
.expect("seeded checkpoint");
let path = write_paginated_rollout(
home.path(),
"2025-01-03T13-00-00",
@@ -103,7 +111,7 @@ async fn stops_at_newest_usable_compaction_and_keeps_companions() {
user_message("older history"),
compacted("older compaction", Some(Vec::new())),
user_message("superseded suffix"),
compacted("latest compaction", Some(Vec::new())),
checkpoint.clone(),
world_state.clone(),
turn_context.clone(),
],
@@ -125,7 +133,7 @@ async fn stops_at_newest_usable_compaction_and_keeps_companions() {
serde_json::to_value(context.items).expect("serialize context"),
serde_json::to_value(vec![
RolloutItem::SessionMeta(Box::new(session_meta)),
compacted("latest compaction", Some(Vec::new())),
checkpoint,
world_state,
turn_context,
])
@@ -272,6 +280,7 @@ async fn fork_version_stops_before_older_segments_once_resolved() {
unreachable!();
};
compacted.resume_metadata = Some(codex_rollout::CompactionResumeMetadata {
next_item_ordinal: None,
rejected_audio_file_ids: Default::default(),
dynamic_tools: None,
multi_agent_version: Some(MultiAgentVersion::V2),
@@ -312,6 +312,7 @@ impl RollbackPlanner {
codex_rollout::RetainedContextEvent::DeliveredAssistantMessage {
message,
acceptance_order,
..
},
) => {
// Nested calls have no response item; match them to instructions by acceptance order.