From 723b23efd0c06289e878b294b29f6d7ed98a66e8 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 02:18:00 +0100 Subject: [PATCH] Reinject missing World State fragments on resume (#30152) ## Why World State restores its structured snapshot on resume so unchanged sections do not have to be rendered again. That is safe only when the model-visible fragment represented by the snapshot is still present in retained history. For selected executor skills, the failing selected-capability scenario exposed this state: ```text persisted World State: selected skill catalog is known retained model history: selected skill catalog message is missing next diff: unchanged, so emit nothing ``` The model resumes without being told about the selected skill catalog. ## What changed World State contributions may now optionally describe the concrete model-visible fragment that must remain in retained history. When a persisted snapshot is present: ```text matching retained fragment exists -> trust snapshot, emit nothing matching retained fragment missing -> treat section as absent, render current state once ``` The skills extension uses this for non-empty selected-environment catalogs by matching its exact rendered catalog body. Empty or hidden catalogs do not require a fragment. ## Scope This does not clear or rebuild the whole World State baseline. It does not change skill discovery, cache invalidation, environment availability, or MCP runtime behavior. It only keeps a persisted section snapshot and its retained model context consistent across resume/history reconstruction. ## Coverage A focused World State regression test verifies both sides: - a missing retained fragment is rendered again - a matching retained fragment avoids duplicate injection --- codex-rs/core/src/context/world_state/mod.rs | 43 +++++++++++- .../context/world_state/world_state_tests.rs | 46 +++++++++++++ .../src/contributors/world_state.rs | 21 ++++++ codex-rs/ext/skills/src/world_state.rs | 66 +++++++++++-------- 4 files changed, 146 insertions(+), 30 deletions(-) diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs index 94960a832..ea82ddd11 100644 --- a/codex-rs/core/src/context/world_state/mod.rs +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -23,6 +23,10 @@ trait ErasedWorldStateSection: Send + Sync { fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool; + fn has_retained_fragment_matcher(&self) -> bool; + + fn matches_retained_fragment(&self, role: &str, text: &str) -> bool; + fn render_diff( &self, previous: PreviousSectionState<'_, Value>, @@ -57,6 +61,14 @@ impl ErasedWorldStateSection for S { S::matches_legacy_fragment(role, text) } + fn has_retained_fragment_matcher(&self) -> bool { + false + } + + fn matches_retained_fragment(&self, _role: &str, _text: &str) -> bool { + false + } + fn render_diff( &self, previous: PreviousSectionState<'_, Value>, @@ -99,6 +111,14 @@ impl ErasedWorldStateSection for ExtensionWorldStateSection { self.0.matches_legacy_fragment(role, text) } + fn has_retained_fragment_matcher(&self) -> bool { + self.0.has_retained_fragment_matcher() + } + + fn matches_retained_fragment(&self, role: &str, text: &str) -> bool { + self.0.matches_retained_fragment(role, text) + } + fn render_diff( &self, previous: PreviousSectionState<'_, Value>, @@ -268,7 +288,12 @@ impl WorldState { ) -> Vec> { self.render_with(|id, section| { if let Some(previous) = previous.and_then(|previous| previous.sections.get(id)) { - PreviousSectionState::Known(previous) + if section.has_retained_fragment_matcher() && !has_retained_fragment(items, section) + { + PreviousSectionState::Absent + } else { + PreviousSectionState::Known(previous) + } } else if has_legacy_fragment(items, section) { PreviousSectionState::Unknown } else { @@ -288,6 +313,22 @@ impl WorldState { } } +fn has_retained_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_retained_fragment(role, text) + ) + }) + ) + }) +} + fn has_legacy_fragment(items: &[ResponseItem], section: &dyn ErasedWorldStateSection) -> bool { items.iter().any(|item| { matches!( 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 7ff6f7328..943267242 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 @@ -149,6 +149,52 @@ fn extension_owned_section_uses_its_snapshot_and_renderer() { ); } +#[test] +fn missing_retained_fragment_is_rendered_again() { + let mut world_state = WorldState::default(); + world_state.add_extension_section( + WorldStateSectionContribution::new( + "extension_test", + json!({"body": "current catalog"}), + |previous| match previous { + PreviousWorldStateSection::Absent => Some(RenderedWorldStateFragment::new( + "developer", + ("", ""), + "current catalog", + )), + PreviousWorldStateSection::Unknown | PreviousWorldStateSection::Known(_) => None, + }, + ) + .with_retained_fragment_matcher(|role, text| { + role == "developer" && text.contains("current catalog") + }), + ); + let previous = world_state.snapshot(); + let retained = ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ContentItem::InputText { + text: "current catalog".to_string(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }; + + assert_eq!( + world_state + .render_history_diff(Some(&previous), &[]) + .into_iter() + .map(|fragment| fragment.body()) + .collect::>(), + vec!["current catalog"] + ); + assert!( + world_state + .render_history_diff(Some(&previous), &[retained]) + .is_empty() + ); +} + #[test] fn unreadable_section_snapshot_is_treated_as_unknown() { let mut current = WorldState::default(); diff --git a/codex-rs/ext/extension-api/src/contributors/world_state.rs b/codex-rs/ext/extension-api/src/contributors/world_state.rs index 23f180826..3bafebffa 100644 --- a/codex-rs/ext/extension-api/src/contributors/world_state.rs +++ b/codex-rs/ext/extension-api/src/contributors/world_state.rs @@ -75,6 +75,7 @@ pub struct WorldStateSectionContribution { snapshot: Value, render_diff: Arc, matches_legacy_fragment: Arc, + matches_retained_fragment: Option>, } impl WorldStateSectionContribution { @@ -93,6 +94,7 @@ impl WorldStateSectionContribution { snapshot, render_diff: Arc::new(render_diff), matches_legacy_fragment: Arc::new(|_, _| false), + matches_retained_fragment: None, } } @@ -104,6 +106,15 @@ impl WorldStateSectionContribution { self } + /// Requires a matching model-visible fragment whenever a persisted snapshot is reused. + pub fn with_retained_fragment_matcher( + mut self, + matcher: impl Fn(&str, &str) -> bool + Send + Sync + 'static, + ) -> Self { + self.matches_retained_fragment = Some(Arc::new(matcher)); + self + } + pub fn id(&self) -> &'static str { self.id } @@ -122,4 +133,14 @@ impl WorldStateSectionContribution { pub fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool { (self.matches_legacy_fragment)(role, text) } + + pub fn has_retained_fragment_matcher(&self) -> bool { + self.matches_retained_fragment.is_some() + } + + pub fn matches_retained_fragment(&self, role: &str, text: &str) -> bool { + self.matches_retained_fragment + .as_ref() + .is_some_and(|matcher| matcher(role, text)) + } } diff --git a/codex-rs/ext/skills/src/world_state.rs b/codex-rs/ext/skills/src/world_state.rs index ce36427b0..e23eab413 100644 --- a/codex-rs/ext/skills/src/world_state.rs +++ b/codex-rs/ext/skills/src/world_state.rs @@ -27,36 +27,44 @@ pub(crate) fn executor_skills_world_state_section( "body": body, "includeInstructions": include_instructions, }); + let retained_body = body.clone(); - WorldStateSectionContribution::new(SKILLS_WORLD_STATE_ID, snapshot, move |previous| { - let previous_is_absent = matches!(&previous, PreviousWorldStateSection::Absent); - if let PreviousWorldStateSection::Known(previous) = &previous { - let previous_body = previous.get("body").and_then(serde_json::Value::as_str); - let previous_include_instructions = previous - .get("includeInstructions") - .and_then(serde_json::Value::as_bool); - if previous_body == body.as_deref() - && previous_include_instructions == Some(include_instructions) - { - return None; + let contribution = + WorldStateSectionContribution::new(SKILLS_WORLD_STATE_ID, snapshot, move |previous| { + let previous_is_absent = matches!(&previous, PreviousWorldStateSection::Absent); + if let PreviousWorldStateSection::Known(previous) = &previous { + let previous_body = previous.get("body").and_then(serde_json::Value::as_str); + let previous_include_instructions = previous + .get("includeInstructions") + .and_then(serde_json::Value::as_bool); + if previous_body == body.as_deref() + && previous_include_instructions == Some(include_instructions) + { + return None; + } } - } - let body = match body.as_deref() { - Some(body) => body, - None if previous_is_absent => return None, - None if !include_instructions => HIDDEN_EXECUTOR_SKILLS_BODY, - None => NO_EXECUTOR_SKILLS_BODY, - }; - Some(RenderedWorldStateFragment::new( - "developer", - (SKILLS_INSTRUCTIONS_OPEN_TAG, SKILLS_INSTRUCTIONS_CLOSE_TAG), - body, - )) - }) - .with_legacy_matcher(|role, text| { - role == "developer" - && text.trim_start().starts_with(SKILLS_INSTRUCTIONS_OPEN_TAG) - && text.trim_end().ends_with(SKILLS_INSTRUCTIONS_CLOSE_TAG) - }) + let body = match body.as_deref() { + Some(body) => body, + None if previous_is_absent => return None, + None if !include_instructions => HIDDEN_EXECUTOR_SKILLS_BODY, + None => NO_EXECUTOR_SKILLS_BODY, + }; + Some(RenderedWorldStateFragment::new( + "developer", + (SKILLS_INSTRUCTIONS_OPEN_TAG, SKILLS_INSTRUCTIONS_CLOSE_TAG), + body, + )) + }) + .with_legacy_matcher(|role, text| { + role == "developer" + && text.trim_start().starts_with(SKILLS_INSTRUCTIONS_OPEN_TAG) + && text.trim_end().ends_with(SKILLS_INSTRUCTIONS_CLOSE_TAG) + }); + match retained_body { + Some(body) => contribution.with_retained_fragment_matcher(move |role, text| { + role == "developer" && text.contains(&body) + }), + None => contribution, + } }