diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index e4f780c00f..929b9b875e 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/app-server/tests/suite/v2/git_attribution.rs b/codex-rs/app-server/tests/suite/v2/git_attribution.rs index f8c17225b9..241be5fab3 100644 --- a/codex-rs/app-server/tests/suite/v2/git_attribution.rs +++ b/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() - && state.remove("git_attribution").is_some() + && world_state.state.remove("git_attribution").is_some() { removed_saved_attribution = true; } 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 9a9c66e483..043dd67839 100644 --- a/codex-rs/core/src/context/world_state/environment_tests.rs +++ b/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": { diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs index abd2c1d096..ac90fc9f55 100644 --- a/codex-rs/core/src/context/world_state/mod.rs +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -278,13 +278,24 @@ pub(crate) struct WorldStateSnapshot { sections: BTreeMap, } +impl From<&Map> for WorldStateSnapshot { + fn from(state: &Map) -> Self { + Self { + sections: state + .iter() + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + } + } +} + impl WorldStateSnapshot { - pub(crate) fn into_value(self) -> Value { - Value::Object(self.sections.into_iter().collect()) + pub(crate) fn into_object(self) -> Map { + 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 { + pub(crate) fn merge_patch_from(&self, previous: &Self) -> Option> { 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<()> { - let Value::Object(patch) = patch else { - return serde_json::from_value::(patch.clone()).map(drop); - }; + pub(crate) fn apply_merge_patch(&mut self, patch: &Map) { // 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(()) } } 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 41e6580908..7f6f629791 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 @@ -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 - .apply_merge_patch( - ¤t - .merge_patch_from(&previous) - .expect("changed snapshots should produce a patch"), - ) - .expect("apply world-state merge patch"); + let patch = current + .merge_patch_from(&previous) + .expect("changed snapshots should produce a patch"); + previous.apply_merge_patch(&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); - } } diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index d953d09644..41d4b27f10 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/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) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8d54232fd9..961b46e9f0 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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) = { diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index 59ae3eba7b..f7f0e8d2fe 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/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()) { - Ok(snapshot) => Some(snapshot), - Err(err) => { - tracing::warn!(%err, "failed to restore world-state snapshot"); - None - } - }; + world_state_baseline = Some(WorldStateSnapshot::from(&world_state.state)); } 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; - } + baseline.apply_merge_patch(&world_state.state); } RolloutItem::SessionMeta(_) | RolloutItem::ResponseItem(_) diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index 44f597f500..5b6614154a 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/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::>(); 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"} }))), ], diff --git a/codex-rs/core/tests/suite/agents_md.rs b/codex-rs/core/tests/suite/agents_md.rs index 5e647c0e34..9007591390 100644 --- a/codex-rs/core/tests/suite/agents_md.rs +++ b/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) }) diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 3702838d37..2cbfb09f67 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -747,7 +747,7 @@ async fn deferred_executor_promotes_primary_environment_when_startup_completes() .collect::>>()? .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] - .state - .pointer("/environments/environments/remote/status"), + world_state_items[0].state["environments"].pointer("/environments/remote/status"), Some(&json!("starting")) ); assert_eq!( - world_state_items[2] - .state - .pointer("/environments/environments/remote/status"), + world_state_items[2].state["environments"].pointer("/environments/remote/status"), Some(&json!("available")) ); assert_eq!( - world_state_items[2] - .state - .pointer("/environments/environments/remote/shell"), + world_state_items[2].state["environments"].pointer("/environments/remote/shell"), Some(&json!("zsh")) ); diff --git a/codex-rs/core/tests/suite/spawn_agent_description.rs b/codex-rs/core/tests/suite/spawn_agent_description.rs index d78511193a..8cb49c11e9 100644 --- a/codex-rs/core/tests/suite/spawn_agent_description.rs +++ b/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::, _>>()? .into_iter() .map(|mut line| { - if let RolloutItem::WorldState(world_state) = &mut line.item - && let Some(state) = world_state.state.as_object_mut() - { + if let RolloutItem::WorldState(world_state) = &mut line.item { + 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") diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 2bd7a8ab97..9167b82f78 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/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, } impl WorldStateItem { - pub fn full(state: Value) -> Self { + pub fn full(state: Map) -> Self { Self { full: true, state } } - pub fn patch(state: Value) -> Self { + pub fn patch(state: Map) -> Self { Self { full: false, state } } }