From f317dc8a17d30d8feb2c79add1d9d565be0402bf Mon Sep 17 00:00:00 2001
From: Tamir Duberstein <tamird@openai.com>
Date: Tue, 11 Aug 2026 21:22:41 +0000
Subject: [PATCH] Reduce cloning in world-state patch handling (#38078)

## What changed

- Deserialize typed section snapshots directly from borrowed JSON values.
- Build and apply world-state merge patches in place instead of cloning and converting entire snapshots.
- Preserve the existing snapshot when a top-level patch has an invalid shape.

## Testing

- Extend the world-state merge-patch test to cover invalid top-level patches and verify that they leave the snapshot unchanged.

GitOrigin-RevId: 6c2212587b025a8dc3878008a0466d4526eee204
---
 codex-rs/core/src/context/world_state/mod.rs  | 50 ++++++++++++++++---
 .../context/world_state/world_state_tests.rs  |  6 +++
 2 files changed, 48 insertions(+), 8 deletions(-)

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<S: WorldStateSection> ErasedWorldStateSection for S {
         let typed_snapshot;
         let previous = match previous {
             PreviousSectionState::Known(previous) => {
-                match serde_json::from_value::<S::Snapshot>(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<Value> {
-        let previous = Value::Object(previous.sections.clone().into_iter().collect());
-        let current = Value::Object(self.sections.clone().into_iter().collect());
-        create_merge_patch(&previous, &current)
+        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::<Self>(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<Value> {
 }
 
 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(&current), 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);
+    }
 }