agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

[3/3] core: replay persisted world state (#29837)

Browse Source 
## Why

Persisting `WorldState` snapshots and patches is only useful if resume
and fork restore that exact comparison baseline. Rebuilding it from
`TurnContextItem` loses section state and can either repeat or suppress
model-visible updates.

This is the third PR in the WorldState persistence stack, built on
#29835.

## What

- Replay full WorldState snapshots and RFC 7386 patches through the
existing rollout reconstruction segments.
- Discard state from rolled-back turns and treat compaction as a
baseline reset.
- Hydrate `ContextManager` from the reconstructed snapshot on resume and
fork.
- Remove the synthetic `TurnContextItem` to WorldState conversion path.
- Leave legacy or malformed rollouts without a baseline so the next
update safely emits a full snapshot.

## Testing

- `just test -p codex-core world_state`
- `just test -p codex-core rollout_reconstruction_tests`
- `just fix -p codex-core`
- `just test -p codex-core` *(the changed tests passed; the full run
also hit unrelated existing/test-environment failures, primarily a
missing `test_stdio_server` binary)*
This commit is contained in:
sayan-oai
2026-06-24 20:32:08 -07:00
committed by GitHub
Unverified
parent df1199fddb
commit a74771340d
9 changed files with 243 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/context/world_state/environment.rs
-48
View File
@@ -5,10 +5,6 @@ use crate::context::environment_context::NetworkContext;
use crate::context::environment_context::push_xml_escaped_text;
use crate::environment_selection::TurnEnvironmentSnapshot;
use crate::session::turn_context::TurnContext;
use codex_exec_server::LOCAL_ENVIRONMENT_ID;
use codex_protocol::protocol::TurnContextItem;
use codex_protocol::protocol::TurnContextNetworkItem;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_path_uri::PathUri;
use serde::Deserialize;
use serde::Serialize;
@@ -43,29 +39,6 @@ impl EnvironmentsState {
}
}
pub(crate) fn from_turn_context_item(turn_context_item: &TurnContextItem) -> Self {
Self {
environments: [(
LOCAL_ENVIRONMENT_ID.to_string(),
EnvironmentState {
cwd: PathUri::from_abs_path(&turn_context_item.cwd),
status: EnvironmentStatus::Available,
shell: None,
},
)]
.into_iter()
.collect(),
current_date: turn_context_item.current_date.clone(),
timezone: turn_context_item.timezone.clone(),
network: network_from_turn_context_item(turn_context_item),
filesystem: Some(FileSystemContext::from_permission_profile(
&turn_context_item.permission_profile(),
&workspace_roots_from_turn_context_item(turn_context_item),
)),
subagents: None,
}
}
pub(crate) fn with_subagents(mut self, subagents: String) -> Self {
if !subagents.is_empty() {
self.subagents = Some(subagents);
@@ -405,27 +378,6 @@ fn network_from_turn_context(turn_context: &TurnContext) -> Option<NetworkContex
))
}
fn network_from_turn_context_item(turn_context_item: &TurnContextItem) -> Option<NetworkContext> {
let TurnContextNetworkItem {
allowed_domains,
denied_domains,
} = turn_context_item.network.as_ref()?;
Some(NetworkContext::new(
allowed_domains.clone(),
denied_domains.clone(),
))
}
fn workspace_roots_from_turn_context_item(
turn_context_item: &TurnContextItem,
) -> Vec<AbsolutePathBuf> {
if let Some(workspace_roots) = turn_context_item.workspace_roots.as_ref() {
return workspace_roots.clone();
}
vec![turn_context_item.cwd.clone()]
}
#[cfg(test)]
#[path = "environment_tests.rs"]
mod tests;
codex-rs/core/src/context/world_state/environment_render_tests.rs
+1 -55
View File
@@ -9,9 +9,7 @@ use codex_protocol::permissions::FileSystemSandboxPolicy;
use codex_protocol::permissions::FileSystemSpecialPath;
use codex_protocol::permissions::NetworkSandboxPolicy;
use codex_protocol::permissions::project_roots_glob_pattern;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::TurnContextItem;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_absolute_path::test_support::PathBufExt;
use core_test_support::test_path_buf;
use pretty_assertions::assert_eq;
@@ -211,58 +209,6 @@ fn serialize_environment_context_with_full_filesystem_profile() {
assert_eq!(context.render(), expected);
}
#[test]
fn turn_context_item_filesystem_uses_workspace_roots_instead_of_cwd() {
let repo = test_abs_path("/repo");
let other_repo = test_abs_path("/other-repo");
let repo_private = repo.join("private");
let item = TurnContextItem {
turn_id: None,
cwd: test_abs_path("/not-the-workspace"),
workspace_roots: Some(vec![repo.clone(), other_repo.clone()]),
current_date: None,
timezone: None,
approval_policy: AskForApproval::Never,
sandbox_policy: SandboxPolicy::new_read_only_policy(),
permission_profile: Some(workspace_write_permission_profile_with_private_denials()),
network: None,
file_system_sandbox_policy: None,
model: "gpt-5".to_string(),
comp_hash: None,
personality: None,
collaboration_mode: None,
multi_agent_version: None,
multi_agent_mode: None,
realtime_active: None,
effort: None,
summary: codex_protocol::config_types::ReasoningSummary::Auto,
};
let context = EnvironmentsState::from_turn_context_item(&item).render();
assert!(
context.contains(&format!(
"<root>{}</root><root>{}</root>",
repo.to_string_lossy(),
other_repo.to_string_lossy()
)),
"{context}"
);
assert!(
context.contains(&format!("<path>{}</path>", repo_private.to_string_lossy())),
"{context}"
);
assert!(
!context.contains(
test_abs_path("/not-the-workspace")
.join("private")
.to_string_lossy()
.as_ref()
),
"{context}"
);
}
#[test]
fn serialize_read_only_environment_context() {
let context = environment_state(
codex-rs/core/src/context/world_state/mod.rs
+26
View File
@@ -100,6 +100,13 @@ impl WorldStateSnapshot {
let current = Value::Object(self.sections.clone().into_iter().collect());
create_merge_patch(&previous, &current)
}
pub(crate) fn apply_merge_patch(&mut self, patch: &Value) -> serde_json::Result<()> {
let mut current = self.clone().into_value();
apply_merge_patch_value(&mut current, patch);
*self = serde_json::from_value(current)?;
Ok(())
}
}
impl fmt::Debug for WorldState {
@@ -193,6 +200,25 @@ fn create_merge_patch(previous: &Value, current: &Value) -> Option<Value> {
Some(Value::Object(patch))
}
fn apply_merge_patch_value(target: &mut Value, patch: &Value) {
let Value::Object(patch) = patch else {
target.clone_from(patch);
return;
};
if !target.is_object() {
*target = Value::Object(Map::new());
}
if let Value::Object(target) = target {
for (key, value) in patch {
if value.is_null() {
target.remove(key);
} else {
apply_merge_patch_value(target.entry(key.clone()).or_insert(Value::Null), value);
}
}
}
}
#[cfg(test)]
#[path = "world_state_tests.rs"]
mod tests;
codex-rs/core/src/context/world_state/world_state_tests.rs
+9 -1
View File
@@ -121,7 +121,7 @@ fn duplicate_section_ids_are_rejected() {
#[test]
fn snapshot_merge_patch_changes_and_removes_nested_values() {
let previous = WorldStateSnapshot {
let mut previous = WorldStateSnapshot {
sections: BTreeMap::from([
(
"kept".to_string(),
@@ -144,5 +144,13 @@ fn snapshot_merge_patch_changes_and_removes_nested_values() {
"removed_section": null,
}))
);
previous
.apply_merge_patch(
&current
.merge_patch_from(&previous)
.expect("changed snapshots should produce a patch"),
)
.expect("apply world-state merge patch");
assert_eq!(previous, current);
assert_eq!(current.merge_patch_from(&current), None);
}
codex-rs/core/src/context_manager/history_tests.rs
+26 -4
View File
@@ -1,6 +1,7 @@
use super::*;
use crate::context::world_state::EnvironmentsState;
use crate::context::UserInstructions;
use crate::context::world_state::WorldState;
use crate::context::world_state::WorldStateSection;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use codex_protocol::AgentPath;
@@ -75,13 +76,34 @@ fn create_history_with_items(items: Vec<ResponseItem>) -> ContextManager {
h
}
struct TestWorldStateSection;
impl WorldStateSection for TestWorldStateSection {
const ID: &'static str = "test";
type Snapshot = bool;
fn snapshot(&self) -> Self::Snapshot {
true
}
fn render_diff(
&self,
previous: Option<&Self::Snapshot>,
) -> Option<Box<dyn crate::context::ContextualUserFragment>> {
(previous != Some(&true)).then(|| {
Box::new(UserInstructions {
directory: None,
text: "test".to_string(),
}) as Box<dyn crate::context::ContextualUserFragment>
})
}
}
#[test]
fn world_state_baseline_deduplicates_until_history_is_replaced() {
let world_state = || {
let mut state = WorldState::default();
state.add_section(EnvironmentsState::from_turn_context_item(
&reference_context_item(),
));
state.add_section(TestWorldStateSection);
state
};
let mut history = ContextManager::new();
codex-rs/core/src/session/mod.rs
+1 -5
View File
@@ -249,7 +249,6 @@ use self::turn::realtime_text_for_event;
use self::turn_context::TurnContext;
use self::turn_context::TurnSkillsContext;
use self::world_state::build_world_state_from_environment_snapshot;
use self::world_state::build_world_state_from_turn_context_item;
#[cfg(test)]
mod rollout_reconstruction_tests;
@@ -1385,6 +1384,7 @@ impl Session {
mut history,
previous_turn_settings,
reference_context_item,
world_state_baseline,
window_number,
first_window_id,
previous_window_id,
@@ -1397,10 +1397,6 @@ impl Session {
// will be processed again if the rollout is reconstructed in a future session.
// This meets image resizing requirements without modifying persisted rollouts.
prepare_response_items(&mut history);
let world_state_baseline = reference_context_item
.as_ref()
.map(build_world_state_from_turn_context_item)
.map(|world_state| world_state.snapshot());
{
let mut state = self.state.lock().await;
state.replace_history(history, reference_context_item);
codex-rs/core/src/session/rollout_reconstruction.rs
+54 -2
View File
@@ -1,4 +1,5 @@
use super::*;
use crate::context::world_state::WorldStateSnapshot;
use crate::context_manager::is_user_turn_boundary;
use codex_protocol::protocol::SessionContextWindow;
use uuid::Uuid;
@@ -10,6 +11,7 @@ pub(super) struct RolloutReconstruction {
pub(super) history: Vec<ResponseItem>,
pub(super) previous_turn_settings: Option<PreviousTurnSettings>,
pub(super) reference_context_item: Option<TurnContextItem>,
pub(super) world_state_baseline: Option<WorldStateSnapshot>,
pub(super) window_number: u64,
pub(super) first_window_id: Option<Uuid>,
pub(super) previous_window_id: Option<Uuid>,
@@ -46,6 +48,7 @@ struct ActiveReplaySegment<'a> {
counts_as_user_turn: bool,
previous_turn_settings: Option<PreviousTurnSettings>,
reference_context_item: TurnReferenceContextItem,
world_state_replay: Vec<&'a RolloutItem>,
base_replacement_history: Option<&'a [ResponseItem]>,
window: Option<ReconstructedWindow>,
}
@@ -60,6 +63,7 @@ fn finalize_active_segment<'a>(
base_replacement_history: &mut Option<&'a [ResponseItem]>,
previous_turn_settings: &mut Option<PreviousTurnSettings>,
reference_context_item: &mut TurnReferenceContextItem,
world_state_replay: &mut Vec<&'a RolloutItem>,
window: &mut Option<ReconstructedWindow>,
pending_rollback_turns: &mut usize,
) {
@@ -73,6 +77,8 @@ fn finalize_active_segment<'a>(
return;
}
world_state_replay.extend(active_segment.world_state_replay);
// A surviving replacement-history checkpoint is a complete history base. Once we
// know the newest surviving one, older rollout items do not affect rebuilt history.
if base_replacement_history.is_none()
@@ -133,6 +139,7 @@ impl Session {
let mut base_replacement_history: Option<&[ResponseItem]> = None;
let mut previous_turn_settings = None;
let mut reference_context_item = TurnReferenceContextItem::NeverSet;
let mut world_state_replay = Vec::new();
let mut window = None;
// Rollback is "drop the newest N user turns". While scanning in reverse, that becomes
// "skip the next N user-turn segments we finalize".
@@ -149,6 +156,7 @@ impl Session {
RolloutItem::Compacted(compacted) => {
let active_segment =
active_segment.get_or_insert_with(ActiveReplaySegment::default);
active_segment.world_state_replay.push(item);
if active_segment.window.is_none()
&& let Some(window_number) = compacted.window_number
{
@@ -235,6 +243,11 @@ impl Session {
}
}
}
RolloutItem::WorldState(_) => {
let active_segment =
active_segment.get_or_insert_with(ActiveReplaySegment::default);
active_segment.world_state_replay.push(item);
}
RolloutItem::EventMsg(EventMsg::TurnStarted(event)) => {
// `TurnStarted` is the oldest boundary of the active reverse segment.
if active_segment.as_ref().is_some_and(|active_segment| {
@@ -249,6 +262,7 @@ impl Session {
&mut base_replacement_history,
&mut previous_turn_settings,
&mut reference_context_item,
&mut world_state_replay,
&mut window,
&mut pending_rollback_turns,
);
@@ -266,8 +280,7 @@ impl Session {
}
RolloutItem::EventMsg(_)
| RolloutItem::SessionMeta(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::WorldState(_) => {}
| RolloutItem::InterAgentCommunicationMetadata { .. } => {}
}
if base_replacement_history.is_some()
@@ -287,6 +300,7 @@ impl Session {
&mut base_replacement_history,
&mut previous_turn_settings,
&mut reference_context_item,
&mut world_state_replay,
&mut window,
&mut pending_rollback_turns,
);
@@ -370,6 +384,43 @@ impl Session {
reference_context_item
};
// Segments and their contents were collected newest-first; replay the surviving records
// chronologically so compaction resets and merge patches have their original meaning.
world_state_replay.reverse();
let mut world_state_baseline: Option<WorldStateSnapshot> = None;
for item in world_state_replay {
match item {
RolloutItem::Compacted(_) => world_state_baseline = None,
RolloutItem::WorldState(world_state) if world_state.full => {
world_state_baseline = match serde_json::from_value(world_state.state.clone()) {
Ok(snapshot) => Some(snapshot),
Err(err) => {
tracing::warn!(%err, "failed to restore world-state snapshot");
None
}
};
}
RolloutItem::WorldState(world_state) => {
let Some(baseline) = world_state_baseline.as_mut() else {
tracing::warn!("ignored world-state patch without a full snapshot");
continue;
};
if let Err(err) = baseline.apply_merge_patch(&world_state.state) {
tracing::warn!(%err, "failed to apply world-state patch");
world_state_baseline = None;
}
}
RolloutItem::SessionMeta(_)
| RolloutItem::ResponseItem(_)
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::TurnContext(_)
| RolloutItem::EventMsg(_) => {
unreachable!("only world-state replay items are collected")
}
}
}
let window = window.or(initial_window).unwrap_or(ReconstructedWindow {
number: fallback_window_number,
first_id: None,
@@ -380,6 +431,7 @@ impl Session {
history: history.into_raw_items(),
previous_turn_settings,
reference_context_item,
world_state_baseline,
window_number: window.number,
first_window_id: window.first_id,
previous_window_id: window.previous_id,
codex-rs/core/src/session/rollout_reconstruction_tests.rs
+126
View File
@@ -1,5 +1,6 @@
use super::*;
use super::tests::build_world_state_from_turn_context;
use super::tests::make_session_and_context;
use codex_protocol::AgentPath;
use codex_protocol::ThreadId;
@@ -12,7 +13,9 @@ use codex_protocol::protocol::ResumedHistory;
use codex_protocol::protocol::SessionContextWindow;
use codex_protocol::protocol::SessionMeta;
use codex_protocol::protocol::SessionMetaLine;
use codex_protocol::protocol::WorldStateItem;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::path::PathBuf;
use uuid::Uuid;
@@ -59,6 +62,49 @@ fn inter_agent_assistant_message(text: &str) -> ResponseItem {
}
}
fn completed_user_turn_rollout(
turn_context_item: TurnContextItem,
items: Vec<RolloutItem>,
) -> Vec<RolloutItem> {
let turn_id = turn_context_item
.turn_id
.clone()
.expect("turn context should have turn_id");
let mut rollout_items = vec![
RolloutItem::EventMsg(EventMsg::TurnStarted(
codex_protocol::protocol::TurnStartedEvent {
turn_id: turn_id.clone(),
trace_id: None,
started_at: None,
model_context_window: Some(128_000),
collaboration_mode_kind: ModeKind::Default,
},
)),
RolloutItem::EventMsg(EventMsg::UserMessage(
codex_protocol::protocol::UserMessageEvent {
client_id: None,
message: "seed".to_string(),
images: None,
local_images: Vec::new(),
text_elements: Vec::new(),
..Default::default()
},
)),
RolloutItem::TurnContext(turn_context_item),
];
rollout_items.extend(items);
rollout_items.push(RolloutItem::EventMsg(EventMsg::TurnComplete(
codex_protocol::protocol::TurnCompleteEvent {
turn_id,
last_agent_message: None,
completed_at: None,
duration_ms: None,
time_to_first_token_ms: None,
},
)));
rollout_items
}
#[tokio::test]
async fn record_initial_history_reconstructs_typed_inter_agent_message() {
let (session, _turn_context) = make_session_and_context().await;
@@ -86,6 +132,31 @@ async fn record_initial_history_reconstructs_typed_inter_agent_message() {
);
}
#[tokio::test]
async fn record_initial_history_restores_world_state_baseline() {
let (session, turn_context) = make_session_and_context().await;
let world_state = build_world_state_from_turn_context(&session, &turn_context).await;
let rollout_items = completed_user_turn_rollout(
turn_context.to_turn_context_item(),
vec![RolloutItem::WorldState(WorldStateItem::full(
world_state.snapshot().into_value(),
))],
);
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
conversation_id: ThreadId::default(),
history: Arc::new(rollout_items),
rollout_path: Some(PathBuf::from("/tmp/resume.jsonl")),
}))
.await;
session
.record_context_updates_and_set_reference_context_item(&turn_context)
.await;
assert_eq!(session.clone_history().await.raw_items(), &[]);
}
#[tokio::test]
async fn record_initial_history_resumed_bare_turn_context_does_not_hydrate_previous_turn_settings()
{
@@ -115,6 +186,11 @@ async fn record_initial_history_resumed_bare_turn_context_does_not_hydrate_previ
};
let rollout_items = vec![RolloutItem::TurnContext(previous_context_item)];
let reconstructed = session
.reconstruct_history_from_rollout(&turn_context, &rollout_items)
.await;
assert_eq!(reconstructed.world_state_baseline, None);
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
conversation_id: ThreadId::default(),
@@ -251,6 +327,9 @@ async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_com
},
)),
RolloutItem::TurnContext(first_context_item.clone()),
RolloutItem::WorldState(WorldStateItem::full(json!({
"test": {"environment": "first"}
}))),
RolloutItem::ResponseItem(turn_one_user.clone()),
RolloutItem::ResponseItem(turn_one_assistant.clone()),
RolloutItem::EventMsg(EventMsg::TurnComplete(
@@ -282,6 +361,9 @@ async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_com
},
)),
RolloutItem::TurnContext(rolled_back_context_item),
RolloutItem::WorldState(WorldStateItem::patch(json!({
"test": {"environment": "rolled-back"}
}))),
RolloutItem::ResponseItem(turn_two_user),
RolloutItem::ResponseItem(turn_two_assistant),
RolloutItem::EventMsg(EventMsg::TurnComplete(
@@ -320,6 +402,11 @@ async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_com
serde_json::to_value(Some(first_context_item))
.expect("serialize expected reference context item")
);
assert_eq!(
serde_json::to_value(reconstructed.world_state_baseline)
.expect("serialize reconstructed world state"),
json!({"test": {"environment": "first"}})
);
}
#[tokio::test]
@@ -986,6 +1073,45 @@ async fn reconstruct_history_prefers_compacted_window_over_session_meta() {
assert_eq!(reconstructed.window_id, Some(compacted_window_id));
}
#[tokio::test]
async fn reconstruct_history_replays_world_state_from_latest_compaction_window() {
let (session, turn_context) = make_session_and_context().await;
let rollout_items = completed_user_turn_rollout(
turn_context.to_turn_context_item(),
vec![
RolloutItem::WorldState(WorldStateItem::full(json!({
"environment": {"status": "old"}
}))),
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: Some(1),
first_window_id: None,
previous_window_id: None,
window_id: None,
}),
RolloutItem::WorldState(WorldStateItem::full(json!({
"environment": {"status": "starting", "cwd": "/workspace"}
}))),
RolloutItem::WorldState(WorldStateItem::patch(json!({
"environment": {"status": "ready"}
}))),
],
);
let reconstructed = session
.reconstruct_history_from_rollout(&turn_context, &rollout_items)
.await;
assert_eq!(
serde_json::to_value(reconstructed.world_state_baseline)
.expect("serialize reconstructed world state"),
json!({
"environment": {"status": "ready", "cwd": "/workspace"}
})
);
}
#[tokio::test]
async fn reconstruct_history_preserves_legacy_compaction_count_with_session_meta_window() {
let (session, turn_context) = make_session_and_context().await;
codex-rs/core/src/session/world_state.rs
-9
View File
@@ -2,7 +2,6 @@ use super::turn_context::TurnContext;
use crate::context::world_state::EnvironmentsState;
use crate::context::world_state::WorldState;
use crate::environment_selection::TurnEnvironmentSnapshot;
use codex_protocol::protocol::TurnContextItem;
pub(super) fn build_world_state_from_environment_snapshot(
turn_context: &TurnContext,
@@ -18,11 +17,3 @@ pub(super) fn build_world_state_from_environment_snapshot(
}
world_state
}
pub(super) fn build_world_state_from_turn_context_item(
turn_context_item: &TurnContextItem,
) -> WorldState {
let mut world_state = WorldState::default();
world_state.add_section(EnvironmentsState::from_turn_context_item(turn_context_item));
world_state
}