Skip to content

Commit a747713

Browse files
authored
[3/3] core: replay persisted world state (#29837)
## 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)*
1 parent df1199f commit a747713

9 files changed

Lines changed: 243 additions & 124 deletions

File tree

codex-rs/core/src/context/world_state/environment.rs

Lines changed: 0 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,6 @@ use crate::context::environment_context::NetworkContext;
55
use crate::context::environment_context::push_xml_escaped_text;
66
use crate::environment_selection::TurnEnvironmentSnapshot;
77
use crate::session::turn_context::TurnContext;
8-
use codex_exec_server::LOCAL_ENVIRONMENT_ID;
9-
use codex_protocol::protocol::TurnContextItem;
10-
use codex_protocol::protocol::TurnContextNetworkItem;
11-
use codex_utils_absolute_path::AbsolutePathBuf;
128
use codex_utils_path_uri::PathUri;
139
use serde::Deserialize;
1410
use serde::Serialize;
@@ -43,29 +39,6 @@ impl EnvironmentsState {
4339
}
4440
}
4541

46-
pub(crate) fn from_turn_context_item(turn_context_item: &TurnContextItem) -> Self {
47-
Self {
48-
environments: [(
49-
LOCAL_ENVIRONMENT_ID.to_string(),
50-
EnvironmentState {
51-
cwd: PathUri::from_abs_path(&turn_context_item.cwd),
52-
status: EnvironmentStatus::Available,
53-
shell: None,
54-
},
55-
)]
56-
.into_iter()
57-
.collect(),
58-
current_date: turn_context_item.current_date.clone(),
59-
timezone: turn_context_item.timezone.clone(),
60-
network: network_from_turn_context_item(turn_context_item),
61-
filesystem: Some(FileSystemContext::from_permission_profile(
62-
&turn_context_item.permission_profile(),
63-
&workspace_roots_from_turn_context_item(turn_context_item),
64-
)),
65-
subagents: None,
66-
}
67-
}
68-
6942
pub(crate) fn with_subagents(mut self, subagents: String) -> Self {
7043
if !subagents.is_empty() {
7144
self.subagents = Some(subagents);
@@ -405,27 +378,6 @@ fn network_from_turn_context(turn_context: &TurnContext) -> Option<NetworkContex
405378
))
406379
}
407380

408-
fn network_from_turn_context_item(turn_context_item: &TurnContextItem) -> Option<NetworkContext> {
409-
let TurnContextNetworkItem {
410-
allowed_domains,
411-
denied_domains,
412-
} = turn_context_item.network.as_ref()?;
413-
Some(NetworkContext::new(
414-
allowed_domains.clone(),
415-
denied_domains.clone(),
416-
))
417-
}
418-
419-
fn workspace_roots_from_turn_context_item(
420-
turn_context_item: &TurnContextItem,
421-
) -> Vec<AbsolutePathBuf> {
422-
if let Some(workspace_roots) = turn_context_item.workspace_roots.as_ref() {
423-
return workspace_roots.clone();
424-
}
425-
426-
vec![turn_context_item.cwd.clone()]
427-
}
428-
429381
#[cfg(test)]
430382
#[path = "environment_tests.rs"]
431383
mod tests;

codex-rs/core/src/context/world_state/environment_render_tests.rs

Lines changed: 1 addition & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,7 @@ use codex_protocol::permissions::FileSystemSandboxPolicy;
99
use codex_protocol::permissions::FileSystemSpecialPath;
1010
use codex_protocol::permissions::NetworkSandboxPolicy;
1111
use codex_protocol::permissions::project_roots_glob_pattern;
12-
use codex_protocol::protocol::AskForApproval;
13-
use codex_protocol::protocol::SandboxPolicy;
14-
use codex_protocol::protocol::TurnContextItem;
12+
use codex_utils_absolute_path::AbsolutePathBuf;
1513
use codex_utils_absolute_path::test_support::PathBufExt;
1614
use core_test_support::test_path_buf;
1715
use pretty_assertions::assert_eq;
@@ -211,58 +209,6 @@ fn serialize_environment_context_with_full_filesystem_profile() {
211209
assert_eq!(context.render(), expected);
212210
}
213211

214-
#[test]
215-
fn turn_context_item_filesystem_uses_workspace_roots_instead_of_cwd() {
216-
let repo = test_abs_path("/repo");
217-
let other_repo = test_abs_path("/other-repo");
218-
let repo_private = repo.join("private");
219-
let item = TurnContextItem {
220-
turn_id: None,
221-
cwd: test_abs_path("/not-the-workspace"),
222-
workspace_roots: Some(vec![repo.clone(), other_repo.clone()]),
223-
current_date: None,
224-
timezone: None,
225-
approval_policy: AskForApproval::Never,
226-
sandbox_policy: SandboxPolicy::new_read_only_policy(),
227-
permission_profile: Some(workspace_write_permission_profile_with_private_denials()),
228-
network: None,
229-
file_system_sandbox_policy: None,
230-
model: "gpt-5".to_string(),
231-
comp_hash: None,
232-
personality: None,
233-
collaboration_mode: None,
234-
multi_agent_version: None,
235-
multi_agent_mode: None,
236-
realtime_active: None,
237-
effort: None,
238-
summary: codex_protocol::config_types::ReasoningSummary::Auto,
239-
};
240-
241-
let context = EnvironmentsState::from_turn_context_item(&item).render();
242-
243-
assert!(
244-
context.contains(&format!(
245-
"<root>{}</root><root>{}</root>",
246-
repo.to_string_lossy(),
247-
other_repo.to_string_lossy()
248-
)),
249-
"{context}"
250-
);
251-
assert!(
252-
context.contains(&format!("<path>{}</path>", repo_private.to_string_lossy())),
253-
"{context}"
254-
);
255-
assert!(
256-
!context.contains(
257-
test_abs_path("/not-the-workspace")
258-
.join("private")
259-
.to_string_lossy()
260-
.as_ref()
261-
),
262-
"{context}"
263-
);
264-
}
265-
266212
#[test]
267213
fn serialize_read_only_environment_context() {
268214
let context = environment_state(

codex-rs/core/src/context/world_state/mod.rs

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,13 @@ impl WorldStateSnapshot {
100100
let current = Value::Object(self.sections.clone().into_iter().collect());
101101
create_merge_patch(&previous, &current)
102102
}
103+
104+
pub(crate) fn apply_merge_patch(&mut self, patch: &Value) -> serde_json::Result<()> {
105+
let mut current = self.clone().into_value();
106+
apply_merge_patch_value(&mut current, patch);
107+
*self = serde_json::from_value(current)?;
108+
Ok(())
109+
}
103110
}
104111

105112
impl fmt::Debug for WorldState {
@@ -193,6 +200,25 @@ fn create_merge_patch(previous: &Value, current: &Value) -> Option<Value> {
193200
Some(Value::Object(patch))
194201
}
195202

203+
fn apply_merge_patch_value(target: &mut Value, patch: &Value) {
204+
let Value::Object(patch) = patch else {
205+
target.clone_from(patch);
206+
return;
207+
};
208+
if !target.is_object() {
209+
*target = Value::Object(Map::new());
210+
}
211+
if let Value::Object(target) = target {
212+
for (key, value) in patch {
213+
if value.is_null() {
214+
target.remove(key);
215+
} else {
216+
apply_merge_patch_value(target.entry(key.clone()).or_insert(Value::Null), value);
217+
}
218+
}
219+
}
220+
}
221+
196222
#[cfg(test)]
197223
#[path = "world_state_tests.rs"]
198224
mod tests;

codex-rs/core/src/context/world_state/world_state_tests.rs

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ fn duplicate_section_ids_are_rejected() {
121121

122122
#[test]
123123
fn snapshot_merge_patch_changes_and_removes_nested_values() {
124-
let previous = WorldStateSnapshot {
124+
let mut previous = WorldStateSnapshot {
125125
sections: BTreeMap::from([
126126
(
127127
"kept".to_string(),
@@ -144,5 +144,13 @@ fn snapshot_merge_patch_changes_and_removes_nested_values() {
144144
"removed_section": null,
145145
}))
146146
);
147+
previous
148+
.apply_merge_patch(
149+
&current
150+
.merge_patch_from(&previous)
151+
.expect("changed snapshots should produce a patch"),
152+
)
153+
.expect("apply world-state merge patch");
154+
assert_eq!(previous, current);
147155
assert_eq!(current.merge_patch_from(&current), None);
148156
}

codex-rs/core/src/context_manager/history_tests.rs

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use super::*;
2-
use crate::context::world_state::EnvironmentsState;
2+
use crate::context::UserInstructions;
33
use crate::context::world_state::WorldState;
4+
use crate::context::world_state::WorldStateSection;
45
use base64::Engine;
56
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
67
use codex_protocol::AgentPath;
@@ -75,13 +76,34 @@ fn create_history_with_items(items: Vec<ResponseItem>) -> ContextManager {
7576
h
7677
}
7778

79+
struct TestWorldStateSection;
80+
81+
impl WorldStateSection for TestWorldStateSection {
82+
const ID: &'static str = "test";
83+
type Snapshot = bool;
84+
85+
fn snapshot(&self) -> Self::Snapshot {
86+
true
87+
}
88+
89+
fn render_diff(
90+
&self,
91+
previous: Option<&Self::Snapshot>,
92+
) -> Option<Box<dyn crate::context::ContextualUserFragment>> {
93+
(previous != Some(&true)).then(|| {
94+
Box::new(UserInstructions {
95+
directory: None,
96+
text: "test".to_string(),
97+
}) as Box<dyn crate::context::ContextualUserFragment>
98+
})
99+
}
100+
}
101+
78102
#[test]
79103
fn world_state_baseline_deduplicates_until_history_is_replaced() {
80104
let world_state = || {
81105
let mut state = WorldState::default();
82-
state.add_section(EnvironmentsState::from_turn_context_item(
83-
&reference_context_item(),
84-
));
106+
state.add_section(TestWorldStateSection);
85107
state
86108
};
87109
let mut history = ContextManager::new();

codex-rs/core/src/session/mod.rs

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -249,7 +249,6 @@ use self::turn::realtime_text_for_event;
249249
use self::turn_context::TurnContext;
250250
use self::turn_context::TurnSkillsContext;
251251
use self::world_state::build_world_state_from_environment_snapshot;
252-
use self::world_state::build_world_state_from_turn_context_item;
253252
#[cfg(test)]
254253
mod rollout_reconstruction_tests;
255254

@@ -1385,6 +1384,7 @@ impl Session {
13851384
mut history,
13861385
previous_turn_settings,
13871386
reference_context_item,
1387+
world_state_baseline,
13881388
window_number,
13891389
first_window_id,
13901390
previous_window_id,
@@ -1397,10 +1397,6 @@ impl Session {
13971397
// will be processed again if the rollout is reconstructed in a future session.
13981398
// This meets image resizing requirements without modifying persisted rollouts.
13991399
prepare_response_items(&mut history);
1400-
let world_state_baseline = reference_context_item
1401-
.as_ref()
1402-
.map(build_world_state_from_turn_context_item)
1403-
.map(|world_state| world_state.snapshot());
14041400
{
14051401
let mut state = self.state.lock().await;
14061402
state.replace_history(history, reference_context_item);

0 commit comments

Comments
 (0)