diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs index 0ea5824925..abd2c1d096 100644 --- a/codex-rs/core/src/context/world_state/mod.rs +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -23,6 +23,7 @@ use codex_extension_api::WorldStateSectionContribution; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use indexmap::IndexMap; +use serde::Deserialize; use serde::Serialize; use serde::de::DeserializeOwned; use serde_json::Map; @@ -109,7 +110,8 @@ impl ErasedWorldStateSection for S { let typed_snapshot; let previous = match previous { PreviousSectionState::Known(previous) => { - match serde_json::from_value::(previous.clone()) { + // Deserialize the borrowed snapshot without copying its JSON tree. + match S::Snapshot::deserialize(previous) { Ok(previous) => { typed_snapshot = previous; PreviousSectionState::Known(&typed_snapshot) @@ -283,15 +285,41 @@ impl WorldStateSnapshot { /// Returns the RFC 7386 merge patch that advances `previous` to `self`. pub(crate) fn merge_patch_from(&self, previous: &Self) -> Option { - let previous = Value::Object(previous.sections.clone().into_iter().collect()); - let current = Value::Object(self.sections.clone().into_iter().collect()); - create_merge_patch(&previous, ¤t) + let mut patch = Map::new(); + // Emit removals first to preserve insertion-ordered JSON patch output. + for key in previous.sections.keys() { + if !self.sections.contains_key(key) { + patch.insert(key.clone(), Value::Null); + } + } + for (key, current) in &self.sections { + if let Some(previous) = previous.sections.get(key) { + if let Some(value) = create_merge_patch(previous, current) { + patch.insert(key.clone(), value); + } + } else { + patch.insert(key.clone(), current.clone()); + } + } + (!patch.is_empty()).then_some(Value::Object(patch)) } 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)?; + let Value::Object(patch) = patch else { + return serde_json::from_value::(patch.clone()).map(drop); + }; + // Borrow existing keys; only newly inserted sections need owned keys. + for (key, value) in patch { + if value.is_null() { + self.sections.remove(key); + } else if let Some(current) = self.sections.get_mut(key) { + apply_merge_patch_value(current, value); + } else { + let mut current = Value::Null; + apply_merge_patch_value(&mut current, value); + self.sections.insert(key.clone(), current); + } + } Ok(()) } } @@ -477,10 +505,12 @@ fn create_merge_patch(previous: &Value, current: &Value) -> Option { } fn apply_merge_patch_value(target: &mut Value, patch: &Value) { + // Nested patches can replace objects with scalars or arrays. let Value::Object(patch) = patch else { target.clone_from(patch); return; }; + // RFC 7386 replaces non-object values with an object before merging. if !target.is_object() { *target = Value::Object(Map::new()); } @@ -488,8 +518,12 @@ fn apply_merge_patch_value(target: &mut Value, patch: &Value) { for (key, value) in patch { if value.is_null() { target.remove(key); + } else if let Some(current) = target.get_mut(key) { + apply_merge_patch_value(current, value); } else { - apply_merge_patch_value(target.entry(key.clone()).or_insert(Value::Null), value); + let mut current = Value::Null; + apply_merge_patch_value(&mut current, value); + target.insert(key.clone(), current); } } } 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 97e1ead64e..41e6580908 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 @@ -273,4 +273,10 @@ fn snapshot_merge_patch_changes_and_removes_nested_values() { .expect("apply world-state merge patch"); assert_eq!(previous, current); assert_eq!(current.merge_patch_from(¤t), None); + + for invalid in [Value::Null, json!(true), json!([])] { + let original = previous.clone(); + assert!(previous.apply_merge_patch(&invalid).is_err()); + assert_eq!(previous, original); + } }