From ab80d4d484609719621826946cd9e137a3558862 Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Thu, 25 Jun 2026 00:03:52 -0700 Subject: [PATCH] core: reconcile legacy WorldState sections (#29997) ## Why Older rollouts can retain model-visible context for a WorldState section without having a persisted snapshot for that section. Treating the missing snapshot as definitely absent can duplicate old context or fail to tell the model that it was replaced or removed. This provides a generic migration path for sections moving into WorldState, beginning with AGENTS.md. Builds on #29810. ## What changed - distinguish section state that is absent, known from a persisted snapshot, or unknown because matching legacy context remains in history - let WorldState sections identify their own legacy fragments while `ContextManager` owns history reconciliation and baseline persistence - make AGENTS.md emit one conservative replacement or removal update for legacy history, then deduplicate from the newly persisted baseline - preserve existing environment rendering when persisted section data is missing or malformed ## Testing - `just test -p codex-core world_state` - `just test -p codex-core cold_resume_invalidates_deleted_legacy_agents_md_once -- --exact` --- .../core/src/context/world_state/agents_md.rs | 25 ++-- .../context/world_state/agents_md_tests.rs | 28 ++++ .../src/context/world_state/environment.rs | 8 +- .../context/world_state/environment_tests.rs | 7 +- codex-rs/core/src/context/world_state/mod.rs | 121 +++++++++++++++--- .../context/world_state/world_state_tests.rs | 37 +++++- codex-rs/core/src/context_manager/history.rs | 20 +-- .../core/src/context_manager/history_tests.rs | 46 ++++++- codex-rs/core/tests/suite/agents_md.rs | 61 ++++++--- 9 files changed, 284 insertions(+), 69 deletions(-) diff --git a/codex-rs/core/src/context/world_state/agents_md.rs b/codex-rs/core/src/context/world_state/agents_md.rs index f713b2aff..1b9a16801 100644 --- a/codex-rs/core/src/context/world_state/agents_md.rs +++ b/codex-rs/core/src/context/world_state/agents_md.rs @@ -1,3 +1,4 @@ +use super::PreviousSectionState; use super::WorldStateSection; use crate::agents_md::LoadedAgentsMd; use crate::context::ContextualUserFragment; @@ -44,27 +45,35 @@ impl WorldStateSection for AgentsMdState { } } + fn matches_legacy_fragment(role: &str, text: &str) -> bool { + role == "user" && UserInstructions::matches_text(text) + } + fn render_diff( &self, - previous: Option<&Self::Snapshot>, + previous: PreviousSectionState<'_, Self::Snapshot>, ) -> Option> { let current = self.snapshot(); - if previous == Some(¤t) { + if matches!(previous, PreviousSectionState::Known(previous) if previous == ¤t) { return None; } - let previous_instructions = previous.and_then(|state| state.text.as_ref()); - let instructions = match (&self.instructions, previous_instructions) { - (Some(instructions), Some(_)) => UserInstructions { + let previous_may_contain_instructions = match previous { + PreviousSectionState::Known(previous) => previous.text.is_some(), + PreviousSectionState::Unknown => true, + PreviousSectionState::Absent => false, + }; + let instructions = match (&self.instructions, previous_may_contain_instructions) { + (Some(instructions), true) => UserInstructions { directory: instructions.directory.clone(), text: format!("{REPLACEMENT_NOTICE}\n\n{}", instructions.text), }, - (Some(instructions), None) => instructions.clone(), - (None, Some(_)) => UserInstructions { + (Some(instructions), false) => instructions.clone(), + (None, true) => UserInstructions { directory: None, text: REMOVAL_NOTICE.to_string(), }, - (None, None) => return None, + (None, false) => return None, }; Some(Box::new(instructions)) } diff --git a/codex-rs/core/src/context/world_state/agents_md_tests.rs b/codex-rs/core/src/context/world_state/agents_md_tests.rs index f4c171b74..d39fe6c72 100644 --- a/codex-rs/core/src/context/world_state/agents_md_tests.rs +++ b/codex-rs/core/src/context/world_state/agents_md_tests.rs @@ -53,6 +53,34 @@ fn changed_and_removed_state_supersedes_previous_instructions() { ); } +#[test] +fn unknown_previous_state_is_explicitly_superseded() { + let loaded = LoadedAgentsMd::from_text_for_testing("current instructions"); + let current = AgentsMdState::new(Some(&loaded)); + assert_eq!( + vec![user_message( + "# AGENTS.md instructions\n\n\nThese AGENTS.md instructions replace all previously provided AGENTS.md instructions.\n\ncurrent instructions\n", + )], + render_fragments(vec![ + WorldStateSection::render_diff(¤t, PreviousSectionState::Unknown) + .expect("unknown state should be replaced"), + ]), + ); + + assert_eq!( + vec![user_message( + "# AGENTS.md instructions\n\n\nThe previously provided AGENTS.md instructions no longer apply.\n", + )], + render_fragments(vec![ + WorldStateSection::render_diff( + &AgentsMdState::default(), + PreviousSectionState::Unknown, + ) + .expect("unknown state should be removed"), + ]), + ); +} + fn render_fragments(fragments: Vec>) -> Vec { fragments .into_iter() diff --git a/codex-rs/core/src/context/world_state/environment.rs b/codex-rs/core/src/context/world_state/environment.rs index 97ce78f1f..d34124ecf 100644 --- a/codex-rs/core/src/context/world_state/environment.rs +++ b/codex-rs/core/src/context/world_state/environment.rs @@ -1,3 +1,4 @@ +use super::PreviousSectionState; use super::WorldStateSection; use crate::context::ContextualUserFragment; use crate::context::environment_context::FileSystemContext; @@ -95,11 +96,14 @@ impl WorldStateSection for EnvironmentsState { fn render_diff( &self, - previous: Option<&Self::Snapshot>, + previous: PreviousSectionState<'_, Self::Snapshot>, ) -> Option> { let current = self.snapshot(); let empty = EnvironmentsSnapshot::default(); - let previous = previous.unwrap_or(&empty); + let previous = match previous { + PreviousSectionState::Known(previous) => previous, + PreviousSectionState::Absent | PreviousSectionState::Unknown => &empty, + }; let turn_context_values_changed = current.current_date != previous.current_date || current.timezone != previous.timezone || current.network != previous.network 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 e96a6c309..70a156121 100644 --- a/codex-rs/core/src/context/world_state/environment_tests.rs +++ b/codex-rs/core/src/context/world_state/environment_tests.rs @@ -220,7 +220,10 @@ fn single_environment_diff_ignores_unknown_shell() -> Result<()> { assert_eq!( None, - render_fragment(WorldStateSection::render_diff(¤t, Some(&previous))) + render_fragment(WorldStateSection::render_diff( + ¤t, + PreviousSectionState::Known(&previous), + )) ); Ok(()) } @@ -248,7 +251,7 @@ fn removed_legacy_environment_renders_unavailable() -> Result<()> { )), render_fragment(WorldStateSection::render_diff( &EnvironmentsState::default(), - Some(&previous), + PreviousSectionState::Known(&previous), )), ); Ok(()) diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs index 541efbe74..50b996709 100644 --- a/codex-rs/core/src/context/world_state/mod.rs +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -2,6 +2,8 @@ mod agents_md; mod environment; use crate::context::ContextualUserFragment; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; use indexmap::IndexMap; use serde::Serialize; use serde::de::DeserializeOwned; @@ -16,7 +18,12 @@ pub(crate) use environment::EnvironmentsState; trait ErasedWorldStateSection: Send + Sync { fn snapshot(&self) -> Option; - fn render_diff(&self, previous: Option<&Value>) -> Option>; + fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool; + + fn render_diff( + &self, + previous: PreviousSectionState<'_, Value>, + ) -> Option>; } impl ErasedWorldStateSection for S { @@ -43,20 +50,47 @@ impl ErasedWorldStateSection for S { Some(snapshot) } - 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()) + fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool { + S::matches_legacy_fragment(role, text) } + + fn render_diff( + &self, + previous: PreviousSectionState<'_, Value>, + ) -> Option> { + let typed_snapshot; + let previous = match previous { + PreviousSectionState::Known(previous) => { + match serde_json::from_value::(previous.clone()) { + Ok(previous) => { + typed_snapshot = previous; + PreviousSectionState::Known(&typed_snapshot) + } + Err(err) => { + tracing::warn!( + section_id = S::ID, + %err, + "failed to restore world-state section snapshot" + ); + PreviousSectionState::Unknown + } + } + } + PreviousSectionState::Absent => PreviousSectionState::Absent, + PreviousSectionState::Unknown => PreviousSectionState::Unknown, + }; + WorldStateSection::render_diff(self, previous) + } +} + +/// What is known about a section's previously model-visible state. +pub(crate) enum PreviousSectionState<'a, T> { + /// No persisted snapshot or matching fragment exists in retained history. + Absent, + /// Retained history contains the section, but its typed snapshot is unavailable. + Unknown, + /// The exact persisted snapshot is available. + Known(&'a T), } /// A typed portion of the state visible to the model. @@ -65,16 +99,22 @@ impl ErasedWorldStateSection for S { /// 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, and must not serialize -/// to null because merge-patch nulls represent deletion. +/// to null because merge-patch nulls represent deletion. Sections migrated +/// from older context can recognize their previous fragments through +/// `matches_legacy_fragment`. pub(crate) trait WorldStateSection: Send + Sync + 'static { const ID: &'static str; type Snapshot: DeserializeOwned + Serialize; fn snapshot(&self) -> Self::Snapshot; + fn matches_legacy_fragment(_role: &str, _text: &str) -> bool { + false + } + fn render_diff( &self, - previous: Option<&Self::Snapshot>, + previous: PreviousSectionState<'_, Self::Snapshot>, ) -> Option>; } @@ -143,21 +183,66 @@ impl WorldState { } } + /// Renders every section as new, without any known previous state. pub(crate) fn render_full(&self) -> Vec> { - self.render_diff(&WorldStateSnapshot::default()) + self.render_with(|_, _| PreviousSectionState::Absent) } + /// Renders each section against the exact persisted snapshot when available. pub(crate) fn render_diff( &self, previous: &WorldStateSnapshot, + ) -> Vec> { + self.render_with(|id, _| match previous.sections.get(id) { + Some(previous) => PreviousSectionState::Known(previous), + None => PreviousSectionState::Absent, + }) + } + + /// Falls back to retained model history when no exact persisted snapshot is available. + pub(crate) fn render_history_diff( + &self, + previous: Option<&WorldStateSnapshot>, + items: &[ResponseItem], + ) -> Vec> { + self.render_with(|id, section| { + if let Some(previous) = previous.and_then(|previous| previous.sections.get(id)) { + PreviousSectionState::Known(previous) + } else if has_legacy_fragment(items, section) { + PreviousSectionState::Unknown + } else { + PreviousSectionState::Absent + } + }) + } + + fn render_with<'a>( + &self, + mut previous: impl FnMut(&str, &dyn ErasedWorldStateSection) -> PreviousSectionState<'a, Value>, ) -> Vec> { self.sections .iter() - .filter_map(|(id, section)| section.render_diff(previous.sections.get(*id))) + .filter_map(|(id, section)| section.render_diff(previous(id, section.as_ref()))) .collect() } } +fn has_legacy_fragment(items: &[ResponseItem], section: &dyn ErasedWorldStateSection) -> bool { + items.iter().any(|item| { + matches!( + item, + ResponseItem::Message { role, content, .. } + if content.iter().any(|content| { + matches!( + content, + ContentItem::InputText { text } + if section.matches_legacy_fragment(role, text) + ) + }) + ) + }) +} + fn remove_null_object_fields(value: &mut Value) { // RFC 7386 reserves object-valued nulls for deletion, but arrays are replaced whole. match value { 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 a3aec922c..d17304bb7 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 @@ -21,11 +21,15 @@ impl WorldStateSection for TestSection { fn render_diff( &self, - previous: Option<&Self::Snapshot>, + previous: PreviousSectionState<'_, Self::Snapshot>, ) -> Option> { - let previous = previous?; - (self.value != previous.value) - .then(|| Box::new(TestFragment(self.value.clone())) as Box) + match previous { + PreviousSectionState::Known(previous) if self.value != previous.value => { + Some(Box::new(TestFragment(self.value.clone()))) + } + PreviousSectionState::Unknown => Some(Box::new(TestFragment("unknown".to_string()))), + PreviousSectionState::Absent | PreviousSectionState::Known(_) => None, + } } } @@ -59,7 +63,7 @@ impl WorldStateSection for DuplicateTestSection { fn render_diff( &self, - _previous: Option<&Self::Snapshot>, + _previous: PreviousSectionState<'_, Self::Snapshot>, ) -> Option> { None } @@ -106,6 +110,29 @@ fn render_diff_restores_the_typed_section_snapshot() { ); } +#[test] +fn unreadable_section_snapshot_is_treated_as_unknown() { + let mut current = WorldState::default(); + current.add_section(TestSection { + value: "current".to_string(), + optional: None, + array: Vec::new(), + }); + let previous = WorldStateSnapshot { + sections: BTreeMap::from([("test".to_string(), json!({"invalid": true}))]), + }; + + let rendered = current.render_diff(&previous); + + assert_eq!( + vec!["unknown"], + rendered + .into_iter() + .map(|fragment| fragment.body()) + .collect::>() + ); +} + #[test] #[should_panic(expected = "duplicate world-state section ID: test")] fn duplicate_section_ids_are_rejected() { diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index defdbeb38..318c5d4d2 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -90,20 +90,14 @@ impl ContextManager { world_state: &WorldState, ) -> (Vec>, Option) { let snapshot = world_state.snapshot(); - let (fragments, rollout_item) = self.world_state_baseline.as_ref().map_or_else( - || { - ( - world_state.render_full(), - Some(WorldStateItem::full(snapshot.clone().into_value())), - ) - }, + let fragments = + world_state.render_history_diff(self.world_state_baseline.as_ref(), &self.items); + let rollout_item = self.world_state_baseline.as_ref().map_or_else( + || Some(WorldStateItem::full(snapshot.clone().into_value())), |previous| { - ( - world_state.render_diff(previous), - snapshot - .merge_patch_from(previous) - .map(WorldStateItem::patch), - ) + snapshot + .merge_patch_from(previous) + .map(WorldStateItem::patch) }, ); self.world_state_baseline = Some(snapshot); diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 396b495d7..5c320b798 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -86,16 +86,25 @@ impl WorldStateSection for TestWorldStateSection { true } + fn matches_legacy_fragment(role: &str, text: &str) -> bool { + role == "user" && UserInstructions::matches_text(text) + } + fn render_diff( &self, - previous: Option<&Self::Snapshot>, + previous: crate::context::world_state::PreviousSectionState<'_, Self::Snapshot>, ) -> Option> { - (previous != Some(&true)).then(|| { - Box::new(UserInstructions { - directory: None, - text: "test".to_string(), - }) as Box + let text = match previous { + crate::context::world_state::PreviousSectionState::Known(true) => return None, + crate::context::world_state::PreviousSectionState::Unknown => "unknown", + crate::context::world_state::PreviousSectionState::Absent + | crate::context::world_state::PreviousSectionState::Known(false) => "test", + }; + Some(Box::new(UserInstructions { + directory: None, + text: text.to_string(), }) + as Box) } } @@ -123,6 +132,31 @@ fn world_state_baseline_deduplicates_until_history_is_replaced() { assert!(replacement_item.is_some_and(|item| item.full)); } +#[test] +fn world_state_reconciles_matching_legacy_history_once() { + let item = crate::context::ContextualUserFragment::into(UserInstructions { + directory: None, + text: "legacy".to_string(), + }); + let mut history = create_history_with_items(vec![item]); + let mut world_state = WorldState::default(); + world_state.add_section(TestWorldStateSection); + + let (fragments, rollout_item) = history.update_world_state(&world_state); + assert_eq!( + vec!["\n\n\nunknown\n"], + fragments + .into_iter() + .map(|fragment| fragment.body()) + .collect::>() + ); + assert!(rollout_item.is_some_and(|item| item.full)); + + let (fragments, rollout_item) = history.update_world_state(&world_state); + assert!(fragments.is_empty()); + assert_eq!(rollout_item, None); +} + fn user_msg(text: &str) -> ResponseItem { ResponseItem::Message { id: None, diff --git a/codex-rs/core/tests/suite/agents_md.rs b/codex-rs/core/tests/suite/agents_md.rs index 894b76027..605e0ada7 100644 --- a/codex-rs/core/tests/suite/agents_md.rs +++ b/codex-rs/core/tests/suite/agents_md.rs @@ -10,6 +10,8 @@ use codex_home::CodexHomeUserInstructionsProvider; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::Op; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; use codex_utils_absolute_path::AbsolutePathBuf; @@ -31,6 +33,7 @@ use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use serde_json::json; +use std::path::Path; use std::sync::Arc; use std::time::Duration; use tempfile::TempDir; @@ -78,6 +81,32 @@ fn write_global_file( Ok(path.abs()) } +fn remove_agents_md_world_state_section(rollout_path: &Path) -> Result<()> { + let rollout = std::fs::read_to_string(rollout_path)?; + let mut removed_section = false; + let retained = rollout + .lines() + .map(serde_json::from_str::) + .collect::, _>>()? + .into_iter() + .map(|mut line| { + if let RolloutItem::WorldState(world_state) = &mut line.item + && let Some(state) = world_state.state.as_object_mut() + { + removed_section |= state.remove("agents_md").is_some(); + } + serde_json::to_string(&line) + }) + .collect::, _>>()? + .join("\n"); + anyhow::ensure!( + removed_section, + "rollout did not contain a persisted AGENTS.md WorldState section" + ); + std::fs::write(rollout_path, format!("{retained}\n"))?; + Ok(()) +} + fn instruction_fragments(request: &responses::ResponsesRequest) -> Vec { request .message_input_texts("user") @@ -771,7 +800,7 @@ async fn invalid_utf8_global_instructions_are_lossy() -> Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn cold_resume_injects_changed_agents_md_once() -> Result<()> { +async fn cold_resume_invalidates_deleted_legacy_agents_md_once() -> Result<()> { // Set up an initial turn and a later cold-resumed turn against the same rollout. let server = responses::start_mock_server().await; let response_mock = responses::mount_sse_sequence( @@ -821,23 +850,20 @@ async fn cold_resume_injects_changed_agents_md_once() -> Result<()> { }) .await; - // Add a preferred override source, then cold-resume with freshly loaded configuration. - let new_source = write_global_file( - home.as_ref(), - GLOBAL_AGENTS_OVERRIDE_FILENAME, - NEW_GLOBAL_INSTRUCTIONS, - )?; - assert_ne!(old_source, new_source); + // Simulate a rollout written before AGENTS.md had a persisted WorldState section. + remove_agents_md_world_state_section(&rollout_path)?; + + std::fs::remove_file(old_source.as_path())?; let mut resume_builder = test_codex().with_home(Arc::clone(&home)); let resumed = resume_builder .resume(&server, Arc::clone(&home), rollout_path) .await?; - // Assert the API reports the new source while model history replays the old structured prefix. + // Model history still contains the old fragment, but the source no longer exists. assert_eq!( resumed.codex.instruction_sources().await, - vec![PathUri::from_abs_path(&new_source)], - "resume reports sources from the newly loaded config" + Vec::::new(), + "resume reports no deleted instruction source" ); resumed.submit_turn("continue resumed thread").await?; @@ -852,11 +878,16 @@ async fn cold_resume_injects_changed_agents_md_once() -> Result<()> { Some(initial_input.as_slice()), "cold resume should replay the original structured input prefix" ); - assert_instruction_replacement_once( - &requests, - OLD_GLOBAL_INSTRUCTIONS, - NEW_GLOBAL_INSTRUCTIONS, + let initial = expected_provider_only_instruction_fragment(OLD_GLOBAL_INSTRUCTIONS); + let removal = expected_provider_only_instruction_fragment( + "The previously provided AGENTS.md instructions no longer apply.", ); + assert_eq!(instruction_fragments(&requests[0]), vec![initial.clone()]); + assert_eq!( + instruction_fragments(&requests[1]), + vec![initial.clone(), removal.clone()] + ); + assert_eq!(instruction_fragments(&requests[2]), vec![initial, removal]); Ok(()) }