diff --git a/codex-rs/core/src/context/world_state/environment.rs b/codex-rs/core/src/context/world_state/environment.rs index 8056a846fe80..97ce78f1fd9e 100644 --- a/codex-rs/core/src/context/world_state/environment.rs +++ b/codex-rs/core/src/context/world_state/environment.rs @@ -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 Option { - 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 { - 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; diff --git a/codex-rs/core/src/context/world_state/environment_render_tests.rs b/codex-rs/core/src/context/world_state/environment_render_tests.rs index be127de554fd..70122586ed2a 100644 --- a/codex-rs/core/src/context/world_state/environment_render_tests.rs +++ b/codex-rs/core/src/context/world_state/environment_render_tests.rs @@ -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!( - "{}{}", - repo.to_string_lossy(), - other_repo.to_string_lossy() - )), - "{context}" - ); - assert!( - context.contains(&format!("{}", 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( diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs index 477753998002..a6f4e0c5f8b1 100644 --- a/codex-rs/core/src/context/world_state/mod.rs +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -100,6 +100,13 @@ impl WorldStateSnapshot { let current = Value::Object(self.sections.clone().into_iter().collect()); create_merge_patch(&previous, ¤t) } + + 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 { 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; diff --git a/codex-rs/core/src/context/world_state/world_state_tests.rs b/codex-rs/core/src/context/world_state/world_state_tests.rs index 2a48a008b179..a3aec922c6d1 100644 --- a/codex-rs/core/src/context/world_state/world_state_tests.rs +++ b/codex-rs/core/src/context/world_state/world_state_tests.rs @@ -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( + ¤t + .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(¤t), None); } diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index c30b17a2f1ed..396b495d7a9c 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -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) -> 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> { + (previous != Some(&true)).then(|| { + Box::new(UserInstructions { + directory: None, + text: "test".to_string(), + }) as Box + }) + } +} + #[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(); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 6e0d70bc17ed..52ce20c90311 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -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); diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index e78c21f5614a..0b89caaf0575 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -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, pub(super) previous_turn_settings: Option, pub(super) reference_context_item: Option, + pub(super) world_state_baseline: Option, pub(super) window_number: u64, pub(super) first_window_id: Option, pub(super) previous_window_id: Option, @@ -46,6 +48,7 @@ struct ActiveReplaySegment<'a> { counts_as_user_turn: bool, previous_turn_settings: Option, reference_context_item: TurnReferenceContextItem, + world_state_replay: Vec<&'a RolloutItem>, base_replacement_history: Option<&'a [ResponseItem]>, window: Option, } @@ -60,6 +63,7 @@ fn finalize_active_segment<'a>( base_replacement_history: &mut Option<&'a [ResponseItem]>, previous_turn_settings: &mut Option, reference_context_item: &mut TurnReferenceContextItem, + world_state_replay: &mut Vec<&'a RolloutItem>, window: &mut Option, 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 = 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, diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index b875e1f49ef1..7c53c93fc417 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -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, +) -> Vec { + 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; diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index 3d46522b8eb5..c20f3df64666 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -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 -}