Let extensions contribute World State sections

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

@@ -2,6 +2,9 @@ mod agents_md;
 mod environment;
 
 use crate::context::ContextualUserFragment;
+use codex_extension_api::PreviousWorldStateSection;
+use codex_extension_api::RenderedWorldStateFragment;
+use codex_extension_api::WorldStateSectionContribution;
 use codex_protocol::models::ContentItem;
 use codex_protocol::models::ResponseItem;
 use indexmap::IndexMap;
@@ -83,6 +86,54 @@ impl<S: WorldStateSection> ErasedWorldStateSection for S {
     }
 }
 
+struct ExtensionWorldStateSection(WorldStateSectionContribution);
+
+impl ErasedWorldStateSection for ExtensionWorldStateSection {
+    fn snapshot(&self) -> Option<Value> {
+        let mut snapshot = self.0.snapshot().clone();
+        remove_null_object_fields(&mut snapshot);
+        (!snapshot.is_null()).then_some(snapshot)
+    }
+
+    fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool {
+        self.0.matches_legacy_fragment(role, text)
+    }
+
+    fn render_diff(
+        &self,
+        previous: PreviousSectionState<'_, Value>,
+    ) -> Option<Box<dyn ContextualUserFragment>> {
+        let previous = match previous {
+            PreviousSectionState::Absent => PreviousWorldStateSection::Absent,
+            PreviousSectionState::Unknown => PreviousWorldStateSection::Unknown,
+            PreviousSectionState::Known(previous) => PreviousWorldStateSection::Known(previous),
+        };
+        self.0
+            .render_diff(previous)
+            .map(|fragment| Box::new(WorldStateContextFragment(fragment)) as _)
+    }
+}
+
+struct WorldStateContextFragment(RenderedWorldStateFragment);
+
+impl ContextualUserFragment for WorldStateContextFragment {
+    fn role(&self) -> &'static str {
+        self.0.role()
+    }
+
+    fn markers(&self) -> (&'static str, &'static str) {
+        self.0.markers()
+    }
+
+    fn body(&self) -> String {
+        self.0.body().to_string()
+    }
+
+    fn type_markers() -> (&'static str, &'static str) {
+        ("", "")
+    }
+}
+
 /// 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.
@@ -169,6 +220,16 @@ impl WorldState {
         self.sections.insert(id, Box::new(section));
     }
 
+    pub(crate) fn add_extension_section(&mut self, section: WorldStateSectionContribution) {
+        let id = section.id();
+        assert!(
+            !self.sections.contains_key(id),
+            "duplicate world-state section ID: {id}"
+        );
+        self.sections
+            .insert(id, Box::new(ExtensionWorldStateSection(section)));
+    }
+
     pub(crate) fn snapshot(&self) -> WorldStateSnapshot {
         WorldStateSnapshot {
             sections: self

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

@@ -110,6 +110,45 @@ fn render_diff_restores_the_typed_section_snapshot() {
     );
 }
 
+#[test]
+fn extension_owned_section_uses_its_snapshot_and_renderer() {
+    let mut world_state = WorldState::default();
+    world_state.add_extension_section(WorldStateSectionContribution::new(
+        "extension_test",
+        json!({"value": "after", "optional": null}),
+        |previous| match previous {
+            PreviousWorldStateSection::Known(previous)
+                if previous == &json!({"value": "before"}) =>
+            {
+                Some(RenderedWorldStateFragment::new(
+                    "developer",
+                    ("<extension_test>", "</extension_test>"),
+                    "after",
+                ))
+            }
+            PreviousWorldStateSection::Absent
+            | PreviousWorldStateSection::Unknown
+            | PreviousWorldStateSection::Known(_) => None,
+        },
+    ));
+    let previous = WorldStateSnapshot {
+        sections: BTreeMap::from([("extension_test".to_string(), json!({"value": "before"}))]),
+    };
+
+    let rendered = world_state.render_diff(&previous);
+
+    assert_eq!(
+        serde_json::to_value(world_state.snapshot()).expect("serialize world-state snapshot"),
+        json!({"extension_test": {"value": "after"}})
+    );
+    assert_eq!(rendered.len(), 1);
+    assert_eq!(rendered[0].role(), "developer");
+    assert_eq!(
+        rendered[0].render(),
+        "<extension_test>after</extension_test>"
+    );
+}
+
 #[test]
 fn unreadable_section_snapshot_is_treated_as_unknown() {
     let mut current = WorldState::default();

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

@@ -490,7 +490,7 @@ impl Session {
         plugins_manager: Arc<PluginsManager>,
         mcp_manager: Arc<McpManager>,
         extensions: Arc<codex_extension_api::ExtensionRegistry<crate::config::Config>>,
-        thread_extension_init: ExtensionDataInit,
+        mut thread_extension_init: ExtensionDataInit,
         supports_openai_form_elicitation: bool,
         agent_control: AgentControl,
         environment_manager: Arc<EnvironmentManager>,
@@ -556,10 +556,17 @@ impl Session {
             config.current_time_reminder.as_ref(),
             external_time_provider,
         )?;
-        let selected_capability_roots = thread_extension_init
-            .get::<Vec<SelectedCapabilityRoot>>()
-            .map(|roots| roots.as_ref().clone())
-            .unwrap_or_else(|| initial_history.get_selected_capability_roots());
+        let selected_capability_roots =
+            match thread_extension_init.get::<Vec<SelectedCapabilityRoot>>() {
+                Some(roots) => roots.as_ref().clone(),
+                None => {
+                    let roots = initial_history.get_selected_capability_roots();
+                    if !roots.is_empty() {
+                        thread_extension_init.insert(roots.clone());
+                    }
+                    roots
+                }
+            };
         let mcp_thread_init = thread_extension_init.clone();
         let thread_extension_data = codex_extension_api::ExtensionData::new_with_init(
             thread_id.to_string(),

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

@@ -3,6 +3,7 @@ use super::step_context::StepContext;
 use crate::context::world_state::AgentsMdState;
 use crate::context::world_state::EnvironmentsState;
 use crate::context::world_state::WorldState;
+use codex_extension_api::WorldStateContributionInput;
 
 impl Session {
     pub(crate) async fn build_world_state_for_step(
@@ -34,6 +35,22 @@ impl Session {
                 .with_subagents(environment_subagents),
             );
         }
+        let environments = step_context.environments.to_selections();
+        for contributor in self.services.extensions.context_contributors() {
+            for section in contributor
+                .contribute_world_state(WorldStateContributionInput {
+                    thread_id: self.thread_id(),
+                    turn_id: turn_context.sub_id.as_str(),
+                    environments: &environments,
+                    session_store: &self.services.session_extension_data,
+                    thread_store: &self.services.thread_extension_data,
+                    turn_store: turn_context.extension_data.as_ref(),
+                })
+                .await
+            {
+                world_state.add_extension_section(section);
+            }
+        }
         world_state
     }
 }

codex-rs/ext/extension-api/Cargo.toml

@@ -19,6 +19,7 @@ codex-context-fragments = { workspace = true }
 codex-protocol = { workspace = true }
 codex-tools = { workspace = true }
 codex-utils-absolute-path = { workspace = true }
+serde_json = { workspace = true }
 
 [dev-dependencies]
 pretty_assertions = { workspace = true }

codex-rs/ext/extension-api/src/contributors.rs

@@ -18,6 +18,7 @@ mod thread_lifecycle;
 mod tool_lifecycle;
 mod turn_input;
 mod turn_lifecycle;
+mod world_state;
 
 pub use context::TurnContextContributionInput;
 pub use mcp::McpServerContribution;
@@ -39,6 +40,10 @@ pub use turn_lifecycle::TurnAbortInput;
 pub use turn_lifecycle::TurnErrorInput;
 pub use turn_lifecycle::TurnStartInput;
 pub use turn_lifecycle::TurnStopInput;
+pub use world_state::PreviousWorldStateSection;
+pub use world_state::RenderedWorldStateFragment;
+pub use world_state::WorldStateContributionInput;
+pub use world_state::WorldStateSectionContribution;
 
 /// Boxed, sendable future returned by asynchronous extension contributors.
 pub type ExtensionFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
@@ -92,6 +97,17 @@ pub trait ContextContributor: Send + Sync {
             Vec::new()
         })
     }
+
+    fn contribute_world_state<'a>(
+        &'a self,
+        input: WorldStateContributionInput<'a>,
+    ) -> ExtensionFuture<'a, Vec<WorldStateSectionContribution>> {
+        Box::pin(async move {
+            let _self = self;
+            let _input = input;
+            Vec::new()
+        })
+    }
 }
 
 /// Contributor for host-owned thread lifecycle gates.

codex-rs/ext/extension-api/src/contributors/world_state.rs

@@ -0,0 +1,122 @@
+use std::sync::Arc;
+
+use codex_protocol::ThreadId;
+use codex_protocol::protocol::TurnEnvironmentSelection;
+use serde_json::Value;
+
+use crate::ExtensionData;
+
+/// Host state available while an extension contributes one sampling step's World State.
+pub struct WorldStateContributionInput<'a> {
+    pub thread_id: ThreadId,
+    pub turn_id: &'a str,
+    pub environments: &'a [TurnEnvironmentSelection],
+    pub session_store: &'a ExtensionData,
+    pub thread_store: &'a ExtensionData,
+    pub turn_store: &'a ExtensionData,
+}
+
+/// What the harness knows about the previous value of one extension-owned section.
+pub enum PreviousWorldStateSection<'a> {
+    Absent,
+    Unknown,
+    Known(&'a Value),
+}
+
+/// Plain model-visible data rendered by an extension-owned World State section.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct RenderedWorldStateFragment {
+    role: &'static str,
+    markers: (&'static str, &'static str),
+    body: String,
+}
+
+impl RenderedWorldStateFragment {
+    pub fn new(
+        role: &'static str,
+        markers: (&'static str, &'static str),
+        body: impl Into<String>,
+    ) -> Self {
+        Self {
+            role,
+            markers,
+            body: body.into(),
+        }
+    }
+
+    pub fn role(&self) -> &'static str {
+        self.role
+    }
+
+    pub fn markers(&self) -> (&'static str, &'static str) {
+        self.markers
+    }
+
+    pub fn body(&self) -> &str {
+        &self.body
+    }
+}
+
+type RenderDiff = dyn for<'a> Fn(PreviousWorldStateSection<'a>) -> Option<RenderedWorldStateFragment>
+    + Send
+    + Sync;
+type LegacyFragmentMatcher = dyn Fn(&str, &str) -> bool + Send + Sync;
+
+/// One extension-owned World State section captured for a sampling step.
+///
+/// The extension owns the stable ID, comparison snapshot, and diff rendering. The harness owns
+/// persistence and the concrete model-context fragment envelope.
+#[derive(Clone)]
+pub struct WorldStateSectionContribution {
+    id: &'static str,
+    snapshot: Value,
+    render_diff: Arc<RenderDiff>,
+    matches_legacy_fragment: Arc<LegacyFragmentMatcher>,
+}
+
+impl WorldStateSectionContribution {
+    pub fn new(
+        id: &'static str,
+        snapshot: Value,
+        render_diff: impl for<'a> Fn(
+            PreviousWorldStateSection<'a>,
+        ) -> Option<RenderedWorldStateFragment>
+        + Send
+        + Sync
+        + 'static,
+    ) -> Self {
+        Self {
+            id,
+            snapshot,
+            render_diff: Arc::new(render_diff),
+            matches_legacy_fragment: Arc::new(|_, _| false),
+        }
+    }
+
+    pub fn with_legacy_matcher(
+        mut self,
+        matcher: impl Fn(&str, &str) -> bool + Send + Sync + 'static,
+    ) -> Self {
+        self.matches_legacy_fragment = Arc::new(matcher);
+        self
+    }
+
+    pub fn id(&self) -> &'static str {
+        self.id
+    }
+
+    pub fn snapshot(&self) -> &Value {
+        &self.snapshot
+    }
+
+    pub fn render_diff(
+        &self,
+        previous: PreviousWorldStateSection<'_>,
+    ) -> Option<RenderedWorldStateFragment> {
+        (self.render_diff)(previous)
+    }
+
+    pub fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool {
+        (self.matches_legacy_fragment)(role, text)
+    }
+}

codex-rs/ext/extension-api/src/lib.rs

@@ -38,8 +38,10 @@ pub use contributors::ExtensionFuture;
 pub use contributors::McpServerContribution;
 pub use contributors::McpServerContributionContext;
 pub use contributors::McpServerContributor;
+pub use contributors::PreviousWorldStateSection;
 pub use contributors::PromptFragment;
 pub use contributors::PromptSlot;
+pub use contributors::RenderedWorldStateFragment;
 pub use contributors::ThreadIdleInput;
 pub use contributors::ThreadLifecycleContributor;
 pub use contributors::ThreadResumeInput;
@@ -63,6 +65,8 @@ pub use contributors::TurnItemContributor;
 pub use contributors::TurnLifecycleContributor;
 pub use contributors::TurnStartInput;
 pub use contributors::TurnStopInput;
+pub use contributors::WorldStateContributionInput;
+pub use contributors::WorldStateSectionContribution;
 pub use registry::ExtensionRegistry;
 pub use registry::ExtensionRegistryBuilder;
 pub use registry::empty_extension_registry;