[2/3] core: persist world state in rollouts (#29835)

## Why

`WorldState` currently remembers its model-visible diff baseline only in
memory. That leaves no durable source for restoring the exact baseline
after resume, fork, rollback, or compaction.

This is the second PR in the WorldState persistence stack, built on
#29833 and following #29249. It records durable state transitions; the
next PR will replay them during rollout reconstruction.

## What

- Add a `world_state` rollout item containing either a full snapshot or
an RFC 7386 JSON Merge Patch.
- Persist a full snapshot after initial context and after compaction
establishes a new context window.
- Persist non-empty patches when later sampling steps or turns advance
the WorldState baseline.
- Write model-visible history before its matching WorldState record, so
an interrupted write can only cause a safe repeated update on replay.
- Preserve WorldState records for full-history forks while excluding
them from thread previews, metadata, and app-server history
materialization.

Older binaries read rollout lines independently, so they skip the
unknown `world_state` records while retaining the rest of the thread.

## Testing

- `just test -p codex-core
snapshot_merge_patch_changes_and_removes_nested_values`
- `just test -p codex-core
world_state_baseline_deduplicates_until_history_is_replaced`
- `just test -p codex-core
deferred_executor_compaction_preserves_then_updates_environment_once`
- `just test -p codex-protocol`
- `just test -p codex-rollout`
- `just test -p codex-state`
- `just test -p codex-thread-store`
- `just test -p codex-app-server-protocol`
This commit is contained in:
sayan-oai
2026-06-24 20:13:49 -07:00
committed by GitHub
Unverified
parent 6db937275f
commit fa036d39aa
22 changed files with 258 additions and 32 deletions
@@ -388,6 +388,7 @@ impl ThreadHistoryBuilder {
RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SessionMeta(_) => {}
}
}
+1 -1
View File
@@ -63,7 +63,7 @@ fn keep_forked_rollout_item(item: &RolloutItem, preserve_reference_context_item:
// Full-history forks preserve the cached prompt prefix and can keep diffing
// from the parent's durable baseline. Truncated forks drop part of that prompt,
// so they must rebuild context on their first child turn.
RolloutItem::TurnContext(_) => preserve_reference_context_item,
RolloutItem::TurnContext(_) | RolloutItem::WorldState(_) => preserve_reference_context_item,
RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => true,
}
}
+1
View File
@@ -169,6 +169,7 @@ async fn persisted_originator(thread: &CodexThread) -> String {
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::EventMsg(_)
| RolloutItem::Compacted(_)
| RolloutItem::WorldState(_)
| RolloutItem::TurnContext(_) => None,
})
.expect("session metadata should be persisted")
+55 -1
View File
@@ -4,6 +4,7 @@ use crate::context::ContextualUserFragment;
use indexmap::IndexMap;
use serde::Serialize;
use serde::de::DeserializeOwned;
use serde_json::Map;
use serde_json::Value;
use std::collections::BTreeMap;
use std::fmt;
@@ -30,6 +31,13 @@ impl<S: WorldStateSection> ErasedWorldStateSection for S {
}
};
remove_null_object_fields(&mut snapshot);
if snapshot.is_null() {
tracing::error!(
section_id = S::ID,
"world-state section snapshot cannot be null"
);
return None;
}
Some(snapshot)
}
@@ -54,7 +62,8 @@ impl<S: WorldStateSection> ErasedWorldStateSection for S {
/// Implementations own how their current state is rendered relative to an
/// earlier snapshot of the same section. `ID` is persisted in rollouts and
/// must remain stable. `Snapshot` should contain only the comparison data
/// needed to decide what the model must be told next.
/// needed to decide what the model must be told next, and must not serialize
/// to null because merge-patch nulls represent deletion.
pub(crate) trait WorldStateSection: Send + Sync + 'static {
const ID: &'static str;
type Snapshot: DeserializeOwned + Serialize;
@@ -80,6 +89,19 @@ pub(crate) struct WorldStateSnapshot {
sections: BTreeMap<String, Value>,
}
impl WorldStateSnapshot {
pub(crate) fn into_value(self) -> Value {
Value::Object(self.sections.into_iter().collect())
}
/// Returns the RFC 7386 merge patch that advances `previous` to `self`.
pub(crate) fn merge_patch_from(&self, previous: &Self) -> Option<Value> {
let previous = Value::Object(previous.sections.clone().into_iter().collect());
let current = Value::Object(self.sections.clone().into_iter().collect());
create_merge_patch(&previous, &current)
}
}
impl fmt::Debug for WorldState {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("WorldState")
@@ -139,6 +161,38 @@ fn remove_null_object_fields(value: &mut Value) {
}
}
fn create_merge_patch(previous: &Value, current: &Value) -> Option<Value> {
if previous == current {
return None;
}
let Value::Object(current) = current else {
return Some(current.clone());
};
let previous = previous.as_object();
let mut patch = Map::new();
if let Some(previous) = previous {
for key in previous.keys() {
if !current.contains_key(key) {
patch.insert(key.clone(), Value::Null);
}
}
}
for (key, current_value) in current {
let Some(previous_value) = previous.and_then(|previous| previous.get(key)) else {
patch.insert(key.clone(), current_value.clone());
continue;
};
if let Some(value_patch) = create_merge_patch(previous_value, current_value) {
patch.insert(key.clone(), value_patch);
}
}
Some(Value::Object(patch))
}
#[cfg(test)]
#[path = "world_state_tests.rs"]
mod tests;
@@ -118,3 +118,31 @@ fn duplicate_section_ids_are_rejected() {
world_state.add_section(DuplicateTestSection);
}
#[test]
fn snapshot_merge_patch_changes_and_removes_nested_values() {
let previous = WorldStateSnapshot {
sections: BTreeMap::from([
(
"kept".to_string(),
json!({"same": true, "changed": "before", "removed": true}),
),
("removed_section".to_string(), json!({"value": true})),
]),
};
let current = WorldStateSnapshot {
sections: BTreeMap::from([(
"kept".to_string(),
json!({"same": true, "changed": "after"}),
)]),
};
assert_eq!(
current.merge_patch_from(&previous),
Some(json!({
"kept": {"changed": "after", "removed": null},
"removed_section": null,
}))
);
assert_eq!(current.merge_patch_from(&current), None);
}
+20 -6
View File
@@ -20,6 +20,7 @@ use codex_protocol::protocol::InterAgentCommunication;
use codex_protocol::protocol::TokenUsage;
use codex_protocol::protocol::TokenUsageInfo;
use codex_protocol::protocol::TurnContextItem;
use codex_protocol::protocol::WorldStateItem;
use codex_utils_cache::BlockingLruCache;
use codex_utils_cache::sha1_digest;
use codex_utils_output_truncation::TruncationPolicy;
@@ -87,13 +88,26 @@ impl ContextManager {
pub(crate) fn update_world_state(
&mut self,
world_state: &WorldState,
) -> Vec<Box<dyn ContextualUserFragment>> {
let fragments = self.world_state_baseline.as_ref().map_or_else(
|| world_state.render_full(),
|previous| world_state.render_diff(previous),
) -> (Vec<Box<dyn ContextualUserFragment>>, Option<WorldStateItem>) {
let snapshot = world_state.snapshot();
let (fragments, rollout_item) = self.world_state_baseline.as_ref().map_or_else(
|| {
(
world_state.render_full(),
Some(WorldStateItem::full(snapshot.clone().into_value())),
)
},
|previous| {
(
world_state.render_diff(previous),
snapshot
.merge_patch_from(previous)
.map(WorldStateItem::patch),
)
},
);
self.world_state_baseline = Some(world_state.snapshot());
fragments
self.world_state_baseline = Some(snapshot);
(fragments, rollout_item)
}
pub(crate) fn set_world_state_baseline(&mut self, snapshot: WorldStateSnapshot) {
@@ -86,12 +86,19 @@ fn world_state_baseline_deduplicates_until_history_is_replaced() {
};
let mut history = ContextManager::new();
assert_eq!(1, history.update_world_state(&world_state()).len());
assert!(history.update_world_state(&world_state()).is_empty());
let (initial_fragments, initial_item) = history.update_world_state(&world_state());
assert_eq!(1, initial_fragments.len());
assert!(initial_item.is_some_and(|item| item.full));
let (unchanged_fragments, unchanged_item) = history.update_world_state(&world_state());
assert!(unchanged_fragments.is_empty());
assert_eq!(unchanged_item, None);
history.replace(Vec::new());
assert_eq!(1, history.update_world_state(&world_state()).len());
let (replacement_fragments, replacement_item) = history.update_world_state(&world_state());
assert_eq!(1, replacement_fragments.len());
assert!(replacement_item.is_some_and(|item| item.full));
}
fn user_msg(text: &str) -> ResponseItem {
+51 -13
View File
@@ -131,6 +131,7 @@ use codex_protocol::protocol::TurnContextNetworkItem;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_protocol::protocol::TurnEnvironmentSelections;
use codex_protocol::protocol::W3cTraceContext;
use codex_protocol::protocol::WorldStateItem;
use codex_protocol::request_permissions::PermissionGrantScope;
use codex_protocol::request_permissions::RequestPermissionProfile;
use codex_protocol::request_permissions::RequestPermissionsArgs;
@@ -2794,8 +2795,14 @@ impl Session {
self.build_world_state_for_environments(turn_context, &step_context.environments)
.await,
);
// Derive the model update and persisted patch from the same two snapshots.
let previous_snapshot = previous_world_state.snapshot();
let world_state_snapshot = world_state.snapshot();
let world_state_item = world_state_snapshot
.merge_patch_from(&previous_snapshot)
.map(WorldStateItem::patch);
let items = crate::context_manager::updates::merge_contextual_fragments(
world_state.render_diff(&previous_world_state.snapshot()),
world_state.render_diff(&previous_snapshot),
);
if !items.is_empty() {
self.record_conversation_items(turn_context, &items).await;
@@ -2806,7 +2813,12 @@ impl Session {
.lock()
.await
.history
.set_world_state_baseline(world_state.snapshot());
.set_world_state_baseline(world_state_snapshot);
// Record the patch after the context it describes is present in model history.
if let Some(world_state_item) = world_state_item {
self.persist_rollout_items(&[RolloutItem::WorldState(world_state_item)])
.await;
}
world_state
}
@@ -2944,18 +2956,25 @@ impl Session {
replacement_history: Some(items.clone()),
..compacted_item
};
// Compaction starts a new history window, so its WorldState baseline must be full.
let mut world_state_item = None;
{
let mut state = self.state.lock().await;
state.replace_history(items, reference_context_item.clone());
if let Some(world_state) = world_state_baseline {
state
.history
.set_world_state_baseline(world_state.snapshot());
let snapshot = world_state.snapshot();
world_state_item = Some(WorldStateItem::full(snapshot.clone().into_value()));
state.history.set_world_state_baseline(snapshot);
}
}
self.persist_rollout_items(&[RolloutItem::Compacted(compacted_item)])
.await;
// Persist the baseline after the replacement history that established it.
if let Some(world_state_item) = world_state_item {
self.persist_rollout_items(&[RolloutItem::WorldState(world_state_item)])
.await;
}
if let Some(turn_context_item) = reference_context_item {
self.persist_rollout_items(&[RolloutItem::TurnContext(turn_context_item)])
.await;
@@ -3525,30 +3544,38 @@ impl Session {
self.build_world_state_for_environments(turn_context, &turn_context.environments)
.await,
);
let mut context_items = if should_inject_full_context {
// Full initial context resets the baseline; later turns persist only its changes.
let (mut context_items, world_state_item) = if should_inject_full_context {
let context_items = self
.build_initial_context_with_world_state(turn_context, world_state.as_ref())
.await;
let snapshot = world_state.snapshot();
self.state
.lock()
.await
.history
.set_world_state_baseline(world_state.snapshot());
context_items
.set_world_state_baseline(snapshot.clone());
(
context_items,
Some(WorldStateItem::full(snapshot.into_value())),
)
} else {
// Steady-state path: append only built-in context diffs here; turn-scoped extension
// context is added below.
let mut context_items = self
.build_settings_update_items(reference_context_item.as_ref(), turn_context)
.await;
let world_state_items = {
let (world_state_items, world_state_item) = {
let mut state = self.state.lock().await;
crate::context_manager::updates::merge_contextual_fragments(
state.history.update_world_state(world_state.as_ref()),
let (fragments, rollout_item) =
state.history.update_world_state(world_state.as_ref());
(
crate::context_manager::updates::merge_contextual_fragments(fragments),
rollout_item,
)
};
context_items.extend(world_state_items);
context_items
(context_items, world_state_item)
};
if !should_inject_full_context && turn_context_changed {
context_items.extend(
@@ -3556,13 +3583,24 @@ impl Session {
.await,
);
}
if !turn_context_changed && context_items.is_empty() {
// A snapshot can change without producing model-visible or TurnContext updates.
let only_world_state_changed = !turn_context_changed && context_items.is_empty();
if only_world_state_changed && world_state_item.is_none() {
return world_state;
}
if !context_items.is_empty() {
self.record_conversation_items(turn_context, &context_items)
.await;
}
// Persist state only after any model-visible context generated from it.
if let Some(world_state_item) = world_state_item {
self.persist_rollout_items(&[RolloutItem::WorldState(world_state_item)])
.await;
}
// A snapshot-only change does not require a duplicate TurnContext record.
if only_world_state_changed {
return world_state;
}
// Persist one `TurnContextItem` per real user turn so resume/lazy replay can recover the
// latest durable baseline even when this turn emitted no model-visible context diffs.
self.persist_rollout_items(&[RolloutItem::TurnContext(turn_context_item.clone())])
@@ -264,8 +264,10 @@ impl Session {
active_segment.get_or_insert_with(ActiveReplaySegment::default);
active_segment.counts_as_user_turn = true;
}
RolloutItem::InterAgentCommunicationMetadata { .. } => {}
RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => {}
RolloutItem::EventMsg(_)
| RolloutItem::SessionMeta(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::WorldState(_) => {}
}
if base_replacement_history.is_some()
@@ -351,6 +353,7 @@ impl Session {
}
RolloutItem::EventMsg(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SessionMeta(_) => {}
}
}
+2
View File
@@ -2744,6 +2744,7 @@ async fn start_new_context_window_assigns_and_persists_item_ids() {
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::EventMsg(_) => None,
});
assert_eq!(
@@ -2802,6 +2803,7 @@ async fn record_initial_history_assigns_and_persists_id_for_forked_response_item
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::EventMsg(_) => None,
});
assert_eq!(persisted_item_id, Some(live_item_id.as_str()));
+42
View File
@@ -22,6 +22,8 @@ use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::ReviewDecision;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::RolloutLine;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_protocol::request_permissions::PermissionGrantScope;
@@ -796,6 +798,46 @@ async fn deferred_executor_compaction_preserves_then_updates_environment_once()
.expect("the next sampling step should report that the environment is ready");
assert!(starting_index < ready_index);
test.codex.ensure_rollout_materialized().await;
test.codex.flush_rollout().await?;
let rollout_path = test.codex.rollout_path().context("rollout path")?;
let rollout = fs::read_to_string(rollout_path)?;
let world_state_items = rollout
.lines()
.map(serde_json::from_str::<RolloutLine>)
.collect::<serde_json::Result<Vec<_>>>()?
.into_iter()
.filter_map(|line| match line.item {
RolloutItem::WorldState(item) => Some(item),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(
world_state_items
.iter()
.map(|item| item.full)
.collect::<Vec<_>>(),
vec![true, true, false]
);
assert_eq!(
world_state_items[0]
.state
.pointer("/environments/environments/remote/status"),
Some(&json!("starting"))
);
assert_eq!(
world_state_items[2]
.state
.pointer("/environments/environments/remote/status"),
Some(&json!("available"))
);
assert_eq!(
world_state_items[2]
.state
.pointer("/environments/environments/remote/shell"),
Some(&json!("zsh"))
);
Ok(())
}
+1
View File
@@ -415,6 +415,7 @@ mod job {
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::EventMsg(_) => None,
})
.collect::<Vec<_>>();
+21
View File
@@ -2603,6 +2603,7 @@ impl InitialHistory {
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::WorldState(_)
| RolloutItem::EventMsg(_) => None,
})
.and_then(|turn_context| turn_context.multi_agent_mode)
@@ -2938,6 +2939,7 @@ fn multi_agent_version_from_items(
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::WorldState(_)
| RolloutItem::EventMsg(_) => None,
})
})
@@ -3092,9 +3094,28 @@ pub enum RolloutItem {
},
Compacted(CompactedItem),
TurnContext(TurnContextItem),
WorldState(WorldStateItem),
EventMsg(EventMsg),
}
/// Persisted comparison state used to resume model-visible world-state diffing.
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema, TS)]
pub struct WorldStateItem {
/// Full snapshots establish a new baseline; patches update the current baseline.
pub full: bool,
pub state: Value,
}
impl WorldStateItem {
pub fn full(state: Value) -> Self {
Self { full: true, state }
}
pub fn patch(state: Value) -> Self {
Self { full: false, state }
}
}
#[derive(Serialize, Clone, Debug, PartialEq, JsonSchema, TS)]
pub struct CompactedItem {
pub message: String,
+4
View File
@@ -1155,6 +1155,9 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result<HeadTai
RolloutItem::TurnContext(_) => {
// Not included in `head`; skip.
}
RolloutItem::WorldState(_) => {
// Not included in `head`; skip.
}
RolloutItem::Compacted(_) => {
// Not included in `head`; skip.
}
@@ -1217,6 +1220,7 @@ pub async fn read_head_for_summary(path: &Path) -> io::Result<Vec<serde_json::Va
RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::EventMsg(_) => {}
}
}
+2
View File
@@ -72,6 +72,7 @@ pub fn builder_from_items(
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::EventMsg(_) => None,
}) && let Some(builder) = builder_from_session_meta(session_meta, rollout_path)
{
@@ -127,6 +128,7 @@ pub async fn extract_metadata_from_rollout(
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::EventMsg(_) => None,
}),
parse_errors,
@@ -232,6 +232,7 @@ fn rollout_item_type(item: &RolloutItem) -> String {
}
RolloutItem::Compacted(_) => "compacted".to_string(),
RolloutItem::TurnContext(_) => "turn_context".to_string(),
RolloutItem::WorldState(_) => "world_state".to_string(),
RolloutItem::EventMsg(EventMsg::ItemCompleted(event)) => {
format!("event.item_completed.{}", turn_item_type(&event.item))
}
+4 -3
View File
@@ -10,9 +10,10 @@ pub fn is_persisted_rollout_item(item: &RolloutItem) -> bool {
| RolloutItem::InterAgentCommunicationMetadata { .. } => true,
RolloutItem::EventMsg(ev) => should_persist_event_msg(ev),
// Persist Codex executive markers so we can analyze flows (e.g., compaction, API turns).
RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => {
true
}
RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SessionMeta(_) => true,
}
}
+1
View File
@@ -1872,6 +1872,7 @@ async fn resume_candidate_matches_cwd(
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::WorldState(_)
| RolloutItem::EventMsg(_) => None,
})
{
+2 -1
View File
@@ -285,7 +285,8 @@ fn conversation_text_from_item(item: &RolloutItem) -> Option<String> {
| RolloutItem::ResponseItem(_)
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_) => None,
| RolloutItem::Compacted(_)
| RolloutItem::WorldState(_) => None,
}
}
+3 -1
View File
@@ -25,6 +25,7 @@ pub fn apply_rollout_item(
RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. } => {}
RolloutItem::Compacted(_) => {}
RolloutItem::WorldState(_) => {}
}
if metadata.model_provider.is_empty() {
metadata.model_provider = default_provider.to_string();
@@ -42,7 +43,8 @@ pub fn rollout_item_affects_thread_metadata(item: &RolloutItem) -> bool {
| RolloutItem::ResponseItem(_)
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_) => false,
| RolloutItem::Compacted(_)
| RolloutItem::WorldState(_) => false,
}
}
+1
View File
@@ -1237,6 +1237,7 @@ pub(super) fn extract_memory_mode(items: &[RolloutItem]) -> Option<String> {
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::EventMsg(_) => None,
})
}
@@ -286,7 +286,8 @@ impl ThreadMetadataSync {
| RolloutItem::ResponseItem(_)
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_) => {}
| RolloutItem::Compacted(_)
| RolloutItem::WorldState(_) => {}
}
}
Some(update)