mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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.
This commit is contained in:
Generated
+1
@@ -3010,6 +3010,7 @@ dependencies = [
|
||||
"codex-tools",
|
||||
"codex-utils-absolute-path",
|
||||
"pretty_assertions",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user