diff --git a/codex-rs/core/src/context/world_state/environment.rs b/codex-rs/core/src/context/world_state/environment.rs index d6e328466..8056a846f 100644 --- a/codex-rs/core/src/context/world_state/environment.rs +++ b/codex-rs/core/src/context/world_state/environment.rs @@ -10,6 +10,8 @@ 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; use std::collections::BTreeMap; /// Environment values visible to the model. @@ -91,17 +93,49 @@ impl EnvironmentsState { } impl WorldStateSection for EnvironmentsState { - fn render_diff(&self, previous: Option<&Self>) -> Option> { - let empty = Self::default(); + const ID: &'static str = "environments"; + type Snapshot = EnvironmentsSnapshot; + + fn snapshot(&self) -> Self::Snapshot { + EnvironmentsSnapshot { + environments: self + .environments + .iter() + .map(|(id, environment)| { + ( + id.clone(), + EnvironmentSnapshot { + cwd: environment.cwd.inferred_native_path_string(), + status: environment.status, + shell: environment.shell.clone(), + }, + ) + }) + .collect(), + current_date: self.current_date.clone(), + timezone: self.timezone.clone(), + network: self.network.as_ref().map(NetworkContext::render), + filesystem: self.filesystem.as_ref().map(FileSystemContext::render), + subagents: self.subagents.clone(), + } + } + + fn render_diff( + &self, + previous: Option<&Self::Snapshot>, + ) -> Option> { + let current = self.snapshot(); + let empty = EnvironmentsSnapshot::default(); let previous = previous.unwrap_or(&empty); - let turn_context_values_changed = self.current_date != previous.current_date - || self.timezone != previous.timezone - || self.network != previous.network - || self.filesystem != previous.filesystem; + let turn_context_values_changed = current.current_date != previous.current_date + || current.timezone != previous.timezone + || current.network != previous.network + || current.filesystem != previous.filesystem; let mut updates = self .environments .iter() - .filter(|(id, environment)| { + .filter(|(id, _)| { + let environment = ¤t.environments[*id]; previous .environments .get(*id) @@ -269,7 +303,24 @@ struct EnvironmentState { shell: Option, } -impl EnvironmentState { +#[derive(Default, Deserialize, Serialize)] +pub(crate) struct EnvironmentsSnapshot { + environments: BTreeMap, + current_date: Option, + timezone: Option, + network: Option, + filesystem: Option, + subagents: Option, +} + +#[derive(Deserialize, Serialize)] +struct EnvironmentSnapshot { + cwd: String, + status: EnvironmentStatus, + shell: Option, +} + +impl EnvironmentSnapshot { fn has_same_diff_value(&self, other: &Self) -> bool { self.cwd == other.cwd && self.status == other.status @@ -281,7 +332,8 @@ impl EnvironmentState { } } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] enum EnvironmentStatus { Starting, Available, diff --git a/codex-rs/core/src/context/world_state/environment_tests.rs b/codex-rs/core/src/context/world_state/environment_tests.rs index 961ee5ed7..e96a6c309 100644 --- a/codex-rs/core/src/context/world_state/environment_tests.rs +++ b/codex-rs/core/src/context/world_state/environment_tests.rs @@ -8,6 +8,7 @@ use codex_protocol::models::PermissionProfile; use codex_protocol::models::ResponseItem; use codex_protocol::permissions::NetworkSandboxPolicy; use pretty_assertions::assert_eq; +use serde_json::json; #[test] fn renders_full_environment_state() -> Result<()> { @@ -95,7 +96,7 @@ fn renders_only_changed_environments() -> Result<()> { "#, )], - render_fragments(current.render_diff(&previous)), + render_fragments(current.render_diff(&previous.snapshot())), ); Ok(()) } @@ -151,7 +152,42 @@ fn persisted_turn_context_values_render_a_diff() -> Result<()> { "#, )], - render_fragments(current.render_diff(&previous)), + render_fragments(current.render_diff(&previous.snapshot())), + ); + Ok(()) +} + +#[test] +fn persisted_snapshot_uses_model_visible_path_and_context_values() -> Result<()> { + let mut world_state = WorldState::default(); + world_state.add_section(EnvironmentsState { + environments: [( + "remote".to_string(), + available("file:///C:/windows", "powershell")?, + )] + .into_iter() + .collect(), + filesystem: Some(FileSystemContext::from_permission_profile( + &PermissionProfile::Disabled, + &[], + )), + ..Default::default() + }); + + assert_eq!( + serde_json::to_value(world_state.snapshot())?, + json!({ + "environments": { + "environments": { + "remote": { + "cwd": "C:\\windows", + "status": "available", + "shell": "powershell" + } + }, + "filesystem": "" + } + }), ); Ok(()) } @@ -180,6 +216,7 @@ fn single_environment_diff_ignores_unknown_shell() -> Result<()> { .collect(), ..Default::default() }; + let previous = WorldStateSection::snapshot(&previous); assert_eq!( None, @@ -199,6 +236,7 @@ fn removed_legacy_environment_renders_unavailable() -> Result<()> { .collect(), ..Default::default() }; + let previous = WorldStateSection::snapshot(&previous); assert_eq!( Some(user_message( diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs index 06f750ef8..e290e9fad 100644 --- a/codex-rs/core/src/context/world_state/mod.rs +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -2,50 +2,82 @@ mod environment; use crate::context::ContextualUserFragment; use indexmap::IndexMap; -use std::any::Any; -use std::any::TypeId; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::Value; +use std::collections::BTreeMap; use std::fmt; pub(crate) use environment::EnvironmentsState; trait ErasedWorldStateSection: Send + Sync { - fn as_any(&self) -> &dyn Any; + fn snapshot(&self) -> Option; - fn render_diff(&self, previous: Option<&dyn Any>) -> Option>; + fn render_diff(&self, previous: Option<&Value>) -> Option>; } impl ErasedWorldStateSection for S { - fn as_any(&self) -> &dyn Any { - self + fn snapshot(&self) -> Option { + let mut snapshot = match serde_json::to_value(WorldStateSection::snapshot(self)) { + Ok(snapshot) => snapshot, + Err(err) => { + tracing::error!( + section_id = S::ID, + %err, + "failed to serialize world-state section snapshot" + ); + return None; + } + }; + remove_null_object_fields(&mut snapshot); + Some(snapshot) } - fn render_diff(&self, previous: Option<&dyn Any>) -> Option> { - let previous = match previous { - Some(previous) => { - let Some(previous) = previous.downcast_ref::() else { - unreachable!("world-state section type must match its type ID"); - }; - Some(previous) - } - None => None, - }; - WorldStateSection::render_diff(self, previous) + fn render_diff(&self, previous: Option<&Value>) -> Option> { + let previous = previous.and_then(|previous| { + serde_json::from_value::(previous.clone()) + .inspect_err(|err| { + tracing::warn!( + section_id = S::ID, + %err, + "failed to restore world-state section snapshot" + ); + }) + .ok() + }); + WorldStateSection::render_diff(self, previous.as_ref()) } } /// A typed portion of the state visible to the model. /// /// Implementations own how their current state is rendered relative to an -/// earlier value of the same section type. A missing previous value requests -/// the section's complete current representation. -pub(crate) trait WorldStateSection: Any + Send + Sync { - fn render_diff(&self, previous: Option<&Self>) -> Option>; +/// 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. +pub(crate) trait WorldStateSection: Send + Sync + 'static { + const ID: &'static str; + type Snapshot: DeserializeOwned + Serialize; + + fn snapshot(&self) -> Self::Snapshot; + + fn render_diff( + &self, + previous: Option<&Self::Snapshot>, + ) -> Option>; } -/// A snapshot of the model-visible world with one section per concrete type. +/// Live model-visible state, keyed by the same stable section IDs used in rollouts. #[derive(Default)] pub(crate) struct WorldState { - sections: IndexMap>, + sections: IndexMap<&'static str, Box>, +} + +/// Compact comparison state for each model-visible world-state section. +#[derive(Clone, Debug, Default, PartialEq, Serialize, serde::Deserialize)] +#[serde(transparent)] +pub(crate) struct WorldStateSnapshot { + sections: BTreeMap, } impl fmt::Debug for WorldState { @@ -58,23 +90,55 @@ impl fmt::Debug for WorldState { impl WorldState { pub(crate) fn add_section(&mut self, section: S) { - self.sections.insert(TypeId::of::(), Box::new(section)); + let id = S::ID; + assert!( + !self.sections.contains_key(id), + "duplicate world-state section ID: {id}" + ); + self.sections.insert(id, Box::new(section)); + } + + pub(crate) fn snapshot(&self) -> WorldStateSnapshot { + WorldStateSnapshot { + sections: self + .sections + .iter() + .filter_map(|(id, section)| { + section + .snapshot() + .map(|snapshot| ((*id).to_string(), snapshot)) + }) + .collect(), + } } pub(crate) fn render_full(&self) -> Vec> { - self.render_diff(&Self::default()) + self.render_diff(&WorldStateSnapshot::default()) } - pub(crate) fn render_diff(&self, previous: &Self) -> Vec> { + pub(crate) fn render_diff( + &self, + previous: &WorldStateSnapshot, + ) -> Vec> { self.sections .iter() - .filter_map(|(type_id, section)| { - let previous = previous - .sections - .get(type_id) - .map(|section| section.as_any()); - section.render_diff(previous) - }) + .filter_map(|(id, section)| section.render_diff(previous.sections.get(*id))) .collect() } } + +fn remove_null_object_fields(value: &mut Value) { + // RFC 7386 reserves object-valued nulls for deletion, but arrays are replaced whole. + match value { + Value::Object(values) => { + values.retain(|_, value| !value.is_null()); + values.values_mut().for_each(remove_null_object_fields); + } + Value::Array(_) => {} + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {} + } +} + +#[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 new file mode 100644 index 000000000..77736b6c2 --- /dev/null +++ b/codex-rs/core/src/context/world_state/world_state_tests.rs @@ -0,0 +1,120 @@ +use super::*; +use pretty_assertions::assert_eq; +use serde::Deserialize; +use serde::Serialize; +use serde_json::json; + +#[derive(Clone, Deserialize, Serialize)] +struct TestSection { + value: String, + optional: Option, + array: Vec, +} + +impl WorldStateSection for TestSection { + const ID: &'static str = "test"; + type Snapshot = Self; + + fn snapshot(&self) -> Self::Snapshot { + self.clone() + } + + fn render_diff( + &self, + previous: Option<&Self::Snapshot>, + ) -> Option> { + let previous = previous?; + (self.value != previous.value) + .then(|| Box::new(TestFragment(self.value.clone())) as Box) + } +} + +struct TestFragment(String); + +impl ContextualUserFragment for TestFragment { + fn role(&self) -> &'static str { + "user" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ("", "") + } + + fn body(&self) -> String { + self.0.clone() + } +} + +struct DuplicateTestSection; + +impl WorldStateSection for DuplicateTestSection { + const ID: &'static str = "test"; + type Snapshot = (); + + fn snapshot(&self) -> Self::Snapshot {} + + fn render_diff( + &self, + _previous: Option<&Self::Snapshot>, + ) -> Option> { + None + } +} + +#[test] +fn snapshot_uses_stable_section_ids_and_omits_null_fields() { + let mut world_state = WorldState::default(); + world_state.add_section(TestSection { + value: "current".to_string(), + optional: None, + array: vec![json!({"value": null})], + }); + + assert_eq!( + serde_json::to_value(world_state.snapshot()).expect("serialize world-state snapshot"), + json!({"test": {"value": "current", "array": [{"value": null}]}}) + ); +} + +#[test] +fn render_diff_restores_the_typed_section_snapshot() { + let mut previous = WorldState::default(); + previous.add_section(TestSection { + value: "before".to_string(), + optional: None, + array: Vec::new(), + }); + let mut current = WorldState::default(); + current.add_section(TestSection { + value: "after".to_string(), + optional: None, + array: Vec::new(), + }); + + let rendered = current.render_diff(&previous.snapshot()); + + assert_eq!( + vec!["after"], + rendered + .into_iter() + .map(|fragment| fragment.body()) + .collect::>() + ); +} + +#[test] +#[should_panic(expected = "duplicate world-state section ID: test")] +fn duplicate_section_ids_are_rejected() { + let mut world_state = WorldState::default(); + world_state.add_section(TestSection { + value: "current".to_string(), + optional: None, + array: Vec::new(), + }); + + world_state.add_section(DuplicateTestSection); +} diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index daa9e3cd4..0fe72f652 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -1,5 +1,6 @@ use crate::context::ContextualUserFragment; use crate::context::world_state::WorldState; +use crate::context::world_state::WorldStateSnapshot; use crate::context_manager::normalize; use crate::event_mapping::has_non_contextual_dev_message_content; use crate::event_mapping::is_contextual_dev_message_content; @@ -29,7 +30,6 @@ use codex_utils_output_truncation::truncate_function_output_items_with_policy; use codex_utils_output_truncation::truncate_text; use std::num::NonZeroUsize; use std::ops::Deref; -use std::sync::Arc; use std::sync::LazyLock; /// Transcript of thread history @@ -52,7 +52,7 @@ pub(crate) struct ContextManager { /// whose non-diff fragments no longer exist in the surviving history. reference_context_item: Option, /// World state most recently appended to model-visible history. - world_state_baseline: Option>, + world_state_baseline: Option, } impl ContextManager { @@ -86,18 +86,18 @@ impl ContextManager { pub(crate) fn update_world_state( &mut self, - world_state: Arc, + world_state: &WorldState, ) -> Vec> { - let fragments = self.world_state_baseline.as_deref().map_or_else( + let fragments = self.world_state_baseline.as_ref().map_or_else( || world_state.render_full(), |previous| world_state.render_diff(previous), ); - self.world_state_baseline = Some(world_state); + self.world_state_baseline = Some(world_state.snapshot()); fragments } - pub(crate) fn set_world_state_baseline(&mut self, world_state: Arc) { - self.world_state_baseline = Some(world_state); + pub(crate) fn set_world_state_baseline(&mut self, snapshot: WorldStateSnapshot) { + self.world_state_baseline = Some(snapshot); } pub(crate) fn set_token_usage_full(&mut self, context_window: i64) { diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 7e7f7c377..937003efa 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -32,7 +32,6 @@ use image::Luma; use image::Rgba; use pretty_assertions::assert_eq; use regex_lite::Regex; -use std::sync::Arc; const EXEC_FORMAT_MAX_BYTES: usize = 10_000; const EXEC_FORMAT_MAX_TOKENS: usize = 2_500; @@ -83,16 +82,16 @@ fn world_state_baseline_deduplicates_until_history_is_replaced() { state.add_section(EnvironmentsState::from_turn_context_item( &reference_context_item(), )); - Arc::new(state) + state }; let mut history = ContextManager::new(); - assert_eq!(1, history.update_world_state(world_state()).len()); - assert!(history.update_world_state(world_state()).is_empty()); + assert_eq!(1, history.update_world_state(&world_state()).len()); + assert!(history.update_world_state(&world_state()).is_empty()); history.replace(Vec::new()); - assert_eq!(1, history.update_world_state(world_state()).len()); + assert_eq!(1, history.update_world_state(&world_state()).len()); } fn user_msg(text: &str) -> ResponseItem { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8a9055aa2..38a8939b7 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1402,14 +1402,13 @@ impl Session { prepare_response_items(&mut history); let world_state_baseline = reference_context_item .as_ref() - .map(build_world_state_from_turn_context_item); + .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); if let Some(world_state) = world_state_baseline { - state - .history - .set_world_state_baseline(Arc::new(world_state)); + state.history.set_world_state_baseline(world_state); } let fallback_ids = state.auto_compact_window_ids(); let window_id = window_id.unwrap_or(fallback_ids.window_id); @@ -2796,7 +2795,7 @@ impl Session { .await, ); let items = crate::context_manager::updates::merge_contextual_fragments( - world_state.render_diff(previous_world_state.as_ref()), + world_state.render_diff(&previous_world_state.snapshot()), ); if !items.is_empty() { self.record_conversation_items(turn_context, &items).await; @@ -2807,7 +2806,7 @@ impl Session { .lock() .await .history - .set_world_state_baseline(Arc::clone(&world_state)); + .set_world_state_baseline(world_state.snapshot()); world_state } @@ -2949,7 +2948,9 @@ impl Session { 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); + state + .history + .set_world_state_baseline(world_state.snapshot()); } } @@ -3532,7 +3533,7 @@ impl Session { .lock() .await .history - .set_world_state_baseline(Arc::clone(&world_state)); + .set_world_state_baseline(world_state.snapshot()); context_items } else { // Steady-state path: append only built-in context diffs here; turn-scoped extension @@ -3543,7 +3544,7 @@ impl Session { let world_state_items = { let mut state = self.state.lock().await; crate::context_manager::updates::merge_contextual_fragments( - state.history.update_world_state(Arc::clone(&world_state)), + state.history.update_world_state(world_state.as_ref()), ) }; context_items.extend(world_state_items); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 983ed2009..da3b9e295 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -8099,11 +8099,13 @@ async fn record_context_updates_includes_turn_context_fragments_on_steady_state_ }); let mut previous_context_item = turn_context.to_turn_context_item(); previous_context_item.turn_id = Some("previous-turn-id".to_string()); - let world_state = Arc::new(build_world_state_from_turn_context(&session, &turn_context).await); + let world_state = build_world_state_from_turn_context(&session, &turn_context).await; { let mut state = session.state.lock().await; state.set_reference_context_item(Some(previous_context_item)); - state.history.set_world_state_baseline(world_state); + state + .history + .set_world_state_baseline(world_state.snapshot()); } session @@ -8773,12 +8775,13 @@ async fn record_context_updates_and_set_reference_context_item_persists_baseline .with_model(next_model.to_string(), &session.services.models_manager) .await; let previous_context_item = previous_context.to_turn_context_item(); - let world_state = - Arc::new(build_world_state_from_turn_context(&session, &previous_context).await); + let world_state = build_world_state_from_turn_context(&session, &previous_context).await; { let mut state = session.state.lock().await; state.set_reference_context_item(Some(previous_context_item.clone())); - state.history.set_world_state_baseline(world_state); + state + .history + .set_world_state_baseline(world_state.snapshot()); } let rollout_path = attach_thread_persistence(&mut session).await;