Skip to content

Commit c9e6d97

Browse files
authored
Let extensions contribute World State sections (#30100)
## Why #29856 already owns the durable thread intent and exact environment binding. This PR adds only the small missing extension boundary: an extension can contribute one named World State section, while core still owns persistence, diffing, and model-visible fragment types. This lets skills stay in the skills extension instead of moving their runtime into core. ## Shape ```text extension-owned state | | contribute section id + JSON snapshot + renderer v core World State | | compare with the previous snapshot v no message, or one incremental model-visible update ``` The extension API is deliberately small: ```rust fn contribute_world_state(...) -> Vec<WorldStateSectionContribution> ``` Core adapts the rendered result to `ContextualUserFragment`, records the snapshot, and keeps the existing compaction/resume behavior. ## What changes - Adds extension-owned World State section contributions. - Calls those contributors from the existing per-step World State builder. - Restores durable selected capability roots into extension thread state on resume. - Keeps the actual model-context fragment and rollout machinery in core. ## What does not change - No skill or MCP implementation moves out of its extension. - No new file watcher, generation, or RPC. - No generic migration of existing World State sections. - No change to the stable environment-ID assumption from #29856. ## Example ```text step 1 snapshot: skills = [] step 2 snapshot: skills = [executor-demo:deploy] core asks the skills extension to render only that change. ``` ## Stack 1. **This PR:** let extensions contribute World State sections. 2. Project executor skills through the skills extension. 3. Pin one MCP runtime to each model step. 4. Project selected MCP/app/connector metadata by environment availability. 5. One end-to-end integration scenario.
1 parent db541f4 commit c9e6d97

9 files changed

Lines changed: 273 additions & 5 deletions

File tree

codex-rs/Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@ mod agents_md;
22
mod environment;
33

44
use crate::context::ContextualUserFragment;
5+
use codex_extension_api::PreviousWorldStateSection;
6+
use codex_extension_api::RenderedWorldStateFragment;
7+
use codex_extension_api::WorldStateSectionContribution;
58
use codex_protocol::models::ContentItem;
69
use codex_protocol::models::ResponseItem;
710
use indexmap::IndexMap;
@@ -83,6 +86,54 @@ impl<S: WorldStateSection> ErasedWorldStateSection for S {
8386
}
8487
}
8588

89+
struct ExtensionWorldStateSection(WorldStateSectionContribution);
90+
91+
impl ErasedWorldStateSection for ExtensionWorldStateSection {
92+
fn snapshot(&self) -> Option<Value> {
93+
let mut snapshot = self.0.snapshot().clone();
94+
remove_null_object_fields(&mut snapshot);
95+
(!snapshot.is_null()).then_some(snapshot)
96+
}
97+
98+
fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool {
99+
self.0.matches_legacy_fragment(role, text)
100+
}
101+
102+
fn render_diff(
103+
&self,
104+
previous: PreviousSectionState<'_, Value>,
105+
) -> Option<Box<dyn ContextualUserFragment>> {
106+
let previous = match previous {
107+
PreviousSectionState::Absent => PreviousWorldStateSection::Absent,
108+
PreviousSectionState::Unknown => PreviousWorldStateSection::Unknown,
109+
PreviousSectionState::Known(previous) => PreviousWorldStateSection::Known(previous),
110+
};
111+
self.0
112+
.render_diff(previous)
113+
.map(|fragment| Box::new(WorldStateContextFragment(fragment)) as _)
114+
}
115+
}
116+
117+
struct WorldStateContextFragment(RenderedWorldStateFragment);
118+
119+
impl ContextualUserFragment for WorldStateContextFragment {
120+
fn role(&self) -> &'static str {
121+
self.0.role()
122+
}
123+
124+
fn markers(&self) -> (&'static str, &'static str) {
125+
self.0.markers()
126+
}
127+
128+
fn body(&self) -> String {
129+
self.0.body().to_string()
130+
}
131+
132+
fn type_markers() -> (&'static str, &'static str) {
133+
("", "")
134+
}
135+
}
136+
86137
/// What is known about a section's previously model-visible state.
87138
pub(crate) enum PreviousSectionState<'a, T> {
88139
/// No persisted snapshot or matching fragment exists in retained history.
@@ -169,6 +220,16 @@ impl WorldState {
169220
self.sections.insert(id, Box::new(section));
170221
}
171222

223+
pub(crate) fn add_extension_section(&mut self, section: WorldStateSectionContribution) {
224+
let id = section.id();
225+
assert!(
226+
!self.sections.contains_key(id),
227+
"duplicate world-state section ID: {id}"
228+
);
229+
self.sections
230+
.insert(id, Box::new(ExtensionWorldStateSection(section)));
231+
}
232+
172233
pub(crate) fn snapshot(&self) -> WorldStateSnapshot {
173234
WorldStateSnapshot {
174235
sections: self

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,45 @@ fn render_diff_restores_the_typed_section_snapshot() {
110110
);
111111
}
112112

113+
#[test]
114+
fn extension_owned_section_uses_its_snapshot_and_renderer() {
115+
let mut world_state = WorldState::default();
116+
world_state.add_extension_section(WorldStateSectionContribution::new(
117+
"extension_test",
118+
json!({"value": "after", "optional": null}),
119+
|previous| match previous {
120+
PreviousWorldStateSection::Known(previous)
121+
if previous == &json!({"value": "before"}) =>
122+
{
123+
Some(RenderedWorldStateFragment::new(
124+
"developer",
125+
("<extension_test>", "</extension_test>"),
126+
"after",
127+
))
128+
}
129+
PreviousWorldStateSection::Absent
130+
| PreviousWorldStateSection::Unknown
131+
| PreviousWorldStateSection::Known(_) => None,
132+
},
133+
));
134+
let previous = WorldStateSnapshot {
135+
sections: BTreeMap::from([("extension_test".to_string(), json!({"value": "before"}))]),
136+
};
137+
138+
let rendered = world_state.render_diff(&previous);
139+
140+
assert_eq!(
141+
serde_json::to_value(world_state.snapshot()).expect("serialize world-state snapshot"),
142+
json!({"extension_test": {"value": "after"}})
143+
);
144+
assert_eq!(rendered.len(), 1);
145+
assert_eq!(rendered[0].role(), "developer");
146+
assert_eq!(
147+
rendered[0].render(),
148+
"<extension_test>after</extension_test>"
149+
);
150+
}
151+
113152
#[test]
114153
fn unreadable_section_snapshot_is_treated_as_unknown() {
115154
let mut current = WorldState::default();

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

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -490,7 +490,7 @@ impl Session {
490490
plugins_manager: Arc<PluginsManager>,
491491
mcp_manager: Arc<McpManager>,
492492
extensions: Arc<codex_extension_api::ExtensionRegistry<crate::config::Config>>,
493-
thread_extension_init: ExtensionDataInit,
493+
mut thread_extension_init: ExtensionDataInit,
494494
supports_openai_form_elicitation: bool,
495495
agent_control: AgentControl,
496496
environment_manager: Arc<EnvironmentManager>,
@@ -556,10 +556,17 @@ impl Session {
556556
config.current_time_reminder.as_ref(),
557557
external_time_provider,
558558
)?;
559-
let selected_capability_roots = thread_extension_init
560-
.get::<Vec<SelectedCapabilityRoot>>()
561-
.map(|roots| roots.as_ref().clone())
562-
.unwrap_or_else(|| initial_history.get_selected_capability_roots());
559+
let selected_capability_roots =
560+
match thread_extension_init.get::<Vec<SelectedCapabilityRoot>>() {
561+
Some(roots) => roots.as_ref().clone(),
562+
None => {
563+
let roots = initial_history.get_selected_capability_roots();
564+
if !roots.is_empty() {
565+
thread_extension_init.insert(roots.clone());
566+
}
567+
roots
568+
}
569+
};
563570
let mcp_thread_init = thread_extension_init.clone();
564571
let thread_extension_data = codex_extension_api::ExtensionData::new_with_init(
565572
thread_id.to_string(),

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use super::step_context::StepContext;
33
use crate::context::world_state::AgentsMdState;
44
use crate::context::world_state::EnvironmentsState;
55
use crate::context::world_state::WorldState;
6+
use codex_extension_api::WorldStateContributionInput;
67

78
impl Session {
89
pub(crate) async fn build_world_state_for_step(
@@ -34,6 +35,22 @@ impl Session {
3435
.with_subagents(environment_subagents),
3536
);
3637
}
38+
let environments = step_context.environments.to_selections();
39+
for contributor in self.services.extensions.context_contributors() {
40+
for section in contributor
41+
.contribute_world_state(WorldStateContributionInput {
42+
thread_id: self.thread_id(),
43+
turn_id: turn_context.sub_id.as_str(),
44+
environments: &environments,
45+
session_store: &self.services.session_extension_data,
46+
thread_store: &self.services.thread_extension_data,
47+
turn_store: turn_context.extension_data.as_ref(),
48+
})
49+
.await
50+
{
51+
world_state.add_extension_section(section);
52+
}
53+
}
3754
world_state
3855
}
3956
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ codex-context-fragments = { workspace = true }
1919
codex-protocol = { workspace = true }
2020
codex-tools = { workspace = true }
2121
codex-utils-absolute-path = { workspace = true }
22+
serde_json = { workspace = true }
2223

2324
[dev-dependencies]
2425
pretty_assertions = { workspace = true }

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ mod thread_lifecycle;
1818
mod tool_lifecycle;
1919
mod turn_input;
2020
mod turn_lifecycle;
21+
mod world_state;
2122

2223
pub use context::TurnContextContributionInput;
2324
pub use mcp::McpServerContribution;
@@ -39,6 +40,10 @@ pub use turn_lifecycle::TurnAbortInput;
3940
pub use turn_lifecycle::TurnErrorInput;
4041
pub use turn_lifecycle::TurnStartInput;
4142
pub use turn_lifecycle::TurnStopInput;
43+
pub use world_state::PreviousWorldStateSection;
44+
pub use world_state::RenderedWorldStateFragment;
45+
pub use world_state::WorldStateContributionInput;
46+
pub use world_state::WorldStateSectionContribution;
4247

4348
/// Boxed, sendable future returned by asynchronous extension contributors.
4449
pub type ExtensionFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
@@ -92,6 +97,17 @@ pub trait ContextContributor: Send + Sync {
9297
Vec::new()
9398
})
9499
}
100+
101+
fn contribute_world_state<'a>(
102+
&'a self,
103+
input: WorldStateContributionInput<'a>,
104+
) -> ExtensionFuture<'a, Vec<WorldStateSectionContribution>> {
105+
Box::pin(async move {
106+
let _self = self;
107+
let _input = input;
108+
Vec::new()
109+
})
110+
}
95111
}
96112

97113
/// Contributor for host-owned thread lifecycle gates.
Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
use std::sync::Arc;
2+
3+
use codex_protocol::ThreadId;
4+
use codex_protocol::protocol::TurnEnvironmentSelection;
5+
use serde_json::Value;
6+
7+
use crate::ExtensionData;
8+
9+
/// Host state available while an extension contributes one sampling step's World State.
10+
pub struct WorldStateContributionInput<'a> {
11+
pub thread_id: ThreadId,
12+
pub turn_id: &'a str,
13+
pub environments: &'a [TurnEnvironmentSelection],
14+
pub session_store: &'a ExtensionData,
15+
pub thread_store: &'a ExtensionData,
16+
pub turn_store: &'a ExtensionData,
17+
}
18+
19+
/// What the harness knows about the previous value of one extension-owned section.
20+
pub enum PreviousWorldStateSection<'a> {
21+
Absent,
22+
Unknown,
23+
Known(&'a Value),
24+
}
25+
26+
/// Plain model-visible data rendered by an extension-owned World State section.
27+
#[derive(Clone, Debug, PartialEq, Eq)]
28+
pub struct RenderedWorldStateFragment {
29+
role: &'static str,
30+
markers: (&'static str, &'static str),
31+
body: String,
32+
}
33+
34+
impl RenderedWorldStateFragment {
35+
pub fn new(
36+
role: &'static str,
37+
markers: (&'static str, &'static str),
38+
body: impl Into<String>,
39+
) -> Self {
40+
Self {
41+
role,
42+
markers,
43+
body: body.into(),
44+
}
45+
}
46+
47+
pub fn role(&self) -> &'static str {
48+
self.role
49+
}
50+
51+
pub fn markers(&self) -> (&'static str, &'static str) {
52+
self.markers
53+
}
54+
55+
pub fn body(&self) -> &str {
56+
&self.body
57+
}
58+
}
59+
60+
type RenderDiff = dyn for<'a> Fn(PreviousWorldStateSection<'a>) -> Option<RenderedWorldStateFragment>
61+
+ Send
62+
+ Sync;
63+
type LegacyFragmentMatcher = dyn Fn(&str, &str) -> bool + Send + Sync;
64+
65+
/// One extension-owned World State section captured for a sampling step.
66+
///
67+
/// The extension owns the stable ID, comparison snapshot, and diff rendering. The harness owns
68+
/// persistence and the concrete model-context fragment envelope.
69+
#[derive(Clone)]
70+
pub struct WorldStateSectionContribution {
71+
id: &'static str,
72+
snapshot: Value,
73+
render_diff: Arc<RenderDiff>,
74+
matches_legacy_fragment: Arc<LegacyFragmentMatcher>,
75+
}
76+
77+
impl WorldStateSectionContribution {
78+
pub fn new(
79+
id: &'static str,
80+
snapshot: Value,
81+
render_diff: impl for<'a> Fn(
82+
PreviousWorldStateSection<'a>,
83+
) -> Option<RenderedWorldStateFragment>
84+
+ Send
85+
+ Sync
86+
+ 'static,
87+
) -> Self {
88+
Self {
89+
id,
90+
snapshot,
91+
render_diff: Arc::new(render_diff),
92+
matches_legacy_fragment: Arc::new(|_, _| false),
93+
}
94+
}
95+
96+
pub fn with_legacy_matcher(
97+
mut self,
98+
matcher: impl Fn(&str, &str) -> bool + Send + Sync + 'static,
99+
) -> Self {
100+
self.matches_legacy_fragment = Arc::new(matcher);
101+
self
102+
}
103+
104+
pub fn id(&self) -> &'static str {
105+
self.id
106+
}
107+
108+
pub fn snapshot(&self) -> &Value {
109+
&self.snapshot
110+
}
111+
112+
pub fn render_diff(
113+
&self,
114+
previous: PreviousWorldStateSection<'_>,
115+
) -> Option<RenderedWorldStateFragment> {
116+
(self.render_diff)(previous)
117+
}
118+
119+
pub fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool {
120+
(self.matches_legacy_fragment)(role, text)
121+
}
122+
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,10 @@ pub use contributors::ExtensionFuture;
3838
pub use contributors::McpServerContribution;
3939
pub use contributors::McpServerContributionContext;
4040
pub use contributors::McpServerContributor;
41+
pub use contributors::PreviousWorldStateSection;
4142
pub use contributors::PromptFragment;
4243
pub use contributors::PromptSlot;
44+
pub use contributors::RenderedWorldStateFragment;
4345
pub use contributors::ThreadIdleInput;
4446
pub use contributors::ThreadLifecycleContributor;
4547
pub use contributors::ThreadResumeInput;
@@ -63,6 +65,8 @@ pub use contributors::TurnItemContributor;
6365
pub use contributors::TurnLifecycleContributor;
6466
pub use contributors::TurnStartInput;
6567
pub use contributors::TurnStopInput;
68+
pub use contributors::WorldStateContributionInput;
69+
pub use contributors::WorldStateSectionContribution;
6670
pub use registry::ExtensionRegistry;
6771
pub use registry::ExtensionRegistryBuilder;
6872
pub use registry::empty_extension_registry;

0 commit comments

Comments
 (0)