Represent persisted world state as JSON objects (#38274)

## Why

World-state snapshots and merge patches are keyed collections of sections. Allowing
their persisted `state` field to contain any JSON value required replay code to
handle shapes that cannot represent world state.

## What changed

- Type `WorldStateItem.state` and snapshot patches as JSON object maps.
- Reconstruct full snapshots directly from those maps and apply patches without
  fallible arbitrary-value conversions.
- Regenerate the app-server protocol schema for the object-only representation.

GitOrigin-RevId: 3d7662b74f3dd1f3ccc5fdd0200074c13a73cbe5

codex-rs/app-server/tests/suite/v2/git_attribution.rs

@@ -354,8 +354,7 @@ fn replace_attribution_fragment_with_legacy(
                 }
             }
             if let RolloutItem::WorldState(world_state) = &mut line.item
-                && let Some(state) = world_state.state.as_object_mut()
+                && world_state.state.remove("git_attribution").is_some()
-                && state.remove("git_attribution").is_some()
             {
                 removed_saved_attribution = true;
             }

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

@@ -194,7 +194,8 @@ fn changing_primary_environment_updates_model_context_and_persisted_state() -> R
     assert_eq!(
         current_world_state
             .snapshot()
-            .merge_patch_from(&previous_world_state.snapshot()),
+            .merge_patch_from(&previous_world_state.snapshot())
+            .map(serde_json::Value::Object),
         Some(json!({
             "environments": {
                 "environments": {

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

@@ -278,13 +278,24 @@ pub(crate) struct WorldStateSnapshot {
     sections: BTreeMap<String, Value>,
 }
 
+impl From<&Map<String, Value>> for WorldStateSnapshot {
+    fn from(state: &Map<String, Value>) -> Self {
+        Self {
+            sections: state
+                .iter()
+                .map(|(key, value)| (key.clone(), value.clone()))
+                .collect(),
+        }
+    }
+}
+
 impl WorldStateSnapshot {
-    pub(crate) fn into_value(self) -> Value {
+    pub(crate) fn into_object(self) -> Map<String, Value> {
-        Value::Object(self.sections.into_iter().collect())
+        self.sections.into_iter().collect()
     }
 
     /// Returns the RFC 7386 merge patch that advances `previous` to `self`.
-    pub(crate) fn merge_patch_from(&self, previous: &Self) -> Option<Value> {
+    pub(crate) fn merge_patch_from(&self, previous: &Self) -> Option<Map<String, Value>> {
         let mut patch = Map::new();
         // Emit removals first to preserve insertion-ordered JSON patch output.
         for key in previous.sections.keys() {
@@ -301,13 +312,10 @@ impl WorldStateSnapshot {
                 patch.insert(key.clone(), current.clone());
             }
         }
-        (!patch.is_empty()).then_some(Value::Object(patch))
+        (!patch.is_empty()).then_some(patch)
     }
 
-    pub(crate) fn apply_merge_patch(&mut self, patch: &Value) -> serde_json::Result<()> {
+    pub(crate) fn apply_merge_patch(&mut self, patch: &Map<String, Value>) {
-        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() {
@@ -320,7 +328,6 @@ impl WorldStateSnapshot {
                 self.sections.insert(key.clone(), current);
             }
         }
-        Ok(())
     }
 }
 

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

@@ -258,25 +258,16 @@ fn snapshot_merge_patch_changes_and_removes_nested_values() {
     };
 
     assert_eq!(
-        current.merge_patch_from(&previous),
+        current.merge_patch_from(&previous).map(Value::Object),
         Some(json!({
             "kept": {"changed": "after", "removed": null},
             "removed_section": null,
         }))
     );
-    previous
+    let patch = current
-        .apply_merge_patch(
+        .merge_patch_from(&previous)
-            &current
+        .expect("changed snapshots should produce a patch");
-                .merge_patch_from(&previous)
+    previous.apply_merge_patch(&patch);
-                .expect("changed snapshots should produce a patch"),
-        )
-        .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);
-    }
 }

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

@@ -100,7 +100,7 @@ impl ContextManager {
         let fragments =
             world_state.render_history_diff(self.world_state_baseline.as_ref(), self.raw_items());
         let rollout_item = self.world_state_baseline.as_ref().map_or_else(
-            || Some(WorldStateItem::full(snapshot.clone().into_value())),
+            || Some(WorldStateItem::full(snapshot.clone().into_object())),
             |previous| {
                 snapshot
                     .merge_patch_from(previous)

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

@@ -3362,7 +3362,7 @@ impl Session {
             state.replace_annotated_history(items, reference_context_item.clone());
             if let Some(world_state) = world_state_baseline {
                 let snapshot = world_state.snapshot();
-                world_state_item = Some(WorldStateItem::full(snapshot.clone().into_value()));
+                world_state_item = Some(WorldStateItem::full(snapshot.clone().into_object()));
                 state.history.set_world_state_baseline(snapshot);
             }
         }
@@ -3799,7 +3799,7 @@ impl Session {
                 .set_world_state_baseline(snapshot.clone());
             (
                 context_items,
-                Some(WorldStateItem::full(snapshot.into_value())),
+                Some(WorldStateItem::full(snapshot.into_object())),
             )
         } else {
             let (world_state_items, world_state_item) = {

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

@@ -395,23 +395,14 @@ impl Session {
             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()) {
+                    world_state_baseline = Some(WorldStateSnapshot::from(&world_state.state));
-                        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) {
+                    baseline.apply_merge_patch(&world_state.state);
-                        tracing::warn!(%err, "failed to apply world-state patch");
-                        world_state_baseline = None;
-                    }
                 }
                 RolloutItem::SessionMeta(_)
                 | RolloutItem::ResponseItem(_)

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

@@ -21,6 +21,12 @@ use serde_json::json;
 use std::path::PathBuf;
 use uuid::Uuid;
 
+macro_rules! object {
+    ($value:tt) => {
+        serde_json::from_value(json!($value)).unwrap()
+    };
+}
+
 fn user_message(text: &str) -> ResponseItem {
     ResponseItem::Message {
         id: None,
@@ -157,7 +163,7 @@ async fn record_initial_history_restores_world_state_baseline() {
         .map(RolloutItem::ResponseItem)
         .collect::<Vec<_>>();
     world_state_items.push(RolloutItem::WorldState(WorldStateItem::full(
-        world_state.snapshot().into_value(),
+        world_state.snapshot().into_object(),
     )));
     let rollout_items =
         completed_user_turn_rollout(turn_context.to_turn_context_item(), world_state_items);
@@ -355,7 +361,7 @@ async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_com
             },
         )),
         RolloutItem::TurnContext(first_context_item.clone()),
-        RolloutItem::WorldState(WorldStateItem::full(json!({
+        RolloutItem::WorldState(WorldStateItem::full(object!({
             "test": {"environment": "first"}
         }))),
         RolloutItem::ResponseItem(turn_one_user.clone().into()),
@@ -391,7 +397,7 @@ async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_com
             },
         )),
         RolloutItem::TurnContext(rolled_back_context_item),
-        RolloutItem::WorldState(WorldStateItem::patch(json!({
+        RolloutItem::WorldState(WorldStateItem::patch(object!({
             "test": {"environment": "rolled-back"}
         }))),
         RolloutItem::ResponseItem(turn_two_user.into()),
@@ -1131,7 +1137,7 @@ async fn reconstruct_history_replays_world_state_from_latest_compaction_window()
     let rollout_items = completed_user_turn_rollout(
         turn_context.to_turn_context_item(),
         vec![
-            RolloutItem::WorldState(WorldStateItem::full(json!({
+            RolloutItem::WorldState(WorldStateItem::full(object!({
                 "environment": {"status": "old"}
             }))),
             RolloutItem::Compacted(CompactedItem {
@@ -1142,10 +1148,10 @@ async fn reconstruct_history_replays_world_state_from_latest_compaction_window()
                 previous_window_id: None,
                 window_id: None,
             }),
-            RolloutItem::WorldState(WorldStateItem::full(json!({
+            RolloutItem::WorldState(WorldStateItem::full(object!({
                 "environment": {"status": "starting", "cwd": "/workspace"}
             }))),
-            RolloutItem::WorldState(WorldStateItem::patch(json!({
+            RolloutItem::WorldState(WorldStateItem::patch(object!({
                 "environment": {"status": "ready"}
             }))),
         ],

codex-rs/core/tests/suite/agents_md.rs

@@ -90,9 +90,9 @@ fn remove_agents_md_world_state_section(rollout_path: &Path) -> Result<()> {
         .into_iter()
         .map(|mut line| {
             if let RolloutItem::WorldState(world_state) = &mut line.item
-                && let Some(state) = world_state.state.as_object_mut()
+                && world_state.state.remove("agents_md").is_some()
             {
-                removed_section |= state.remove("agents_md").is_some();
+                removed_section = true;
             }
             serde_json::to_string(&line)
         })

codex-rs/core/tests/suite/remote_env.rs

@@ -747,7 +747,7 @@ async fn deferred_executor_promotes_primary_environment_when_startup_completes()
         .collect::<serde_json::Result<Vec<_>>>()?
         .into_iter()
         .filter_map(|line| match line.item {
-            RolloutItem::WorldState(item) if !item.full => Some(item.state),
+            RolloutItem::WorldState(item) if !item.full => Some(Value::Object(item.state)),
             _ => None,
         })
         .find(|patch| {
@@ -1982,21 +1982,15 @@ async fn deferred_executor_compaction_preserves_then_updates_environment_once()
         vec![true, true, false]
     );
     assert_eq!(
-        world_state_items[0]
+        world_state_items[0].state["environments"].pointer("/environments/remote/status"),
-            .state
-            .pointer("/environments/environments/remote/status"),
         Some(&json!("starting"))
     );
     assert_eq!(
-        world_state_items[2]
+        world_state_items[2].state["environments"].pointer("/environments/remote/status"),
-            .state
-            .pointer("/environments/environments/remote/status"),
         Some(&json!("available"))
     );
     assert_eq!(
-        world_state_items[2]
+        world_state_items[2].state["environments"].pointer("/environments/remote/shell"),
-            .state
-            .pointer("/environments/environments/remote/shell"),
         Some(&json!("zsh"))
     );
 

codex-rs/core/tests/suite/spawn_agent_description.rs

@@ -440,9 +440,8 @@ async fn multi_agent_v2_cold_resume_refreshes_legacy_usage_hints_once(
         .collect::<std::result::Result<Vec<_>, _>>()?
         .into_iter()
         .map(|mut line| {
-            if let RolloutItem::WorldState(world_state) = &mut line.item
+            if let RolloutItem::WorldState(world_state) = &mut line.item {
-                && let Some(state) = world_state.state.as_object_mut()
+                let state = &mut world_state.state;
-            {
                 removed_recorded_usage_hint |= state.remove("multi_agent_usage_hint").is_some();
                 if let Some(mode) = state
                     .get_mut("multi_agent_mode")

codex-rs/protocol/src/protocol.rs

@@ -63,6 +63,7 @@ use serde::Deserialize;
 use serde::Deserializer;
 use serde::Serialize;
 use serde::de::Error as _;
+use serde_json::Map;
 use serde_json::Value;
 use serde_with::serde_as;
 use strum_macros::Display;
@@ -2996,15 +2997,15 @@ impl<'de> Deserialize<'de> for SessionMetaLine {
 pub struct WorldStateItem {
     /// Full snapshots establish a new baseline; patches update the current baseline.
     pub full: bool,
-    pub state: Value,
+    pub state: Map<String, Value>,
 }
 
 impl WorldStateItem {
-    pub fn full(state: Value) -> Self {
+    pub fn full(state: Map<String, Value>) -> Self {
         Self { full: true, state }
     }
 
-    pub fn patch(state: Value) -> Self {
+    pub fn patch(state: Map<String, Value>) -> Self {
         Self { full: false, state }
     }
 }