From cfa422c672de706f614fe3171b7da3b2167570eb Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 25 Jun 2026 09:53:17 +0100 Subject: [PATCH] Bind selected capability roots to model steps --- codex-rs/core/src/environment_selection.rs | 74 +++++++++++++++++++ codex-rs/core/src/session/mod.rs | 3 + codex-rs/core/src/session/session.rs | 1 + codex-rs/core/src/session/step_context.rs | 5 ++ codex-rs/core/src/session/tests.rs | 3 + codex-rs/core/src/session/world_state.rs | 4 + codex-rs/core/src/state/service.rs | 4 + codex-rs/core/src/tools/spec_plan_tests.rs | 1 + codex-rs/exec-server/src/lib.rs | 2 + .../exec-server/src/resolved_capability.rs | 47 ++++++++++++ 10 files changed, 144 insertions(+) create mode 100644 codex-rs/exec-server/src/resolved_capability.rs diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index fa86dded55..fcab147725 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -7,6 +7,9 @@ use codex_exec_server::Environment; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecServerError; use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::ResolvedSelectedCapabilityRoot; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; @@ -227,6 +230,27 @@ pub(crate) struct TurnEnvironmentSnapshot { } impl TurnEnvironmentSnapshot { + pub(crate) fn resolve_selected_capability_roots( + &self, + selected_roots: &[SelectedCapabilityRoot], + ) -> Vec { + selected_roots + .iter() + .filter_map(|selected_root| { + let CapabilityRootLocation::Environment { environment_id, .. } = + &selected_root.location; + let environment = self + .turn_environments + .iter() + .find(|environment| environment.environment_id == *environment_id)?; + Some(ResolvedSelectedCapabilityRoot::new( + selected_root.clone(), + Arc::clone(&environment.environment), + )) + }) + .collect() + } + pub(crate) fn primary(&self) -> Option<&TurnEnvironment> { self.turn_environments.first() } @@ -560,6 +584,56 @@ url = "ws://127.0.0.1:8765" assert_eq!(resolved.snapshot().await.to_selections(), vec![local]); } + #[tokio::test] + async fn selected_capability_roots_bind_only_to_their_ready_environment() { + let cwd = PathUri::parse("file:///workspace").expect("cwd URI"); + let primary = Arc::new(Environment::default_for_tests()); + let selected = Arc::new(Environment::default_for_tests()); + let snapshot = TurnEnvironmentSnapshot { + turn_environments: vec![ + TurnEnvironment::new( + "primary".to_string(), + Arc::clone(&primary), + cwd.clone(), + /*shell*/ None, + ), + TurnEnvironment::new( + "selected".to_string(), + Arc::clone(&selected), + cwd, + /*shell*/ None, + ), + ], + starting: Vec::new(), + }; + let plugin_path = PathUri::parse("file:///plugins/demo").expect("plugin URI"); + let roots = [ + SelectedCapabilityRoot { + id: "demo@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "selected".to_string(), + path: plugin_path.clone(), + }, + }, + SelectedCapabilityRoot { + id: "pending@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "pending".to_string(), + path: plugin_path, + }, + }, + ]; + + let resolved = snapshot.resolve_selected_capability_roots(&roots); + + let [resolved] = resolved.as_slice() else { + panic!("only the ready selected environment should resolve"); + }; + assert_eq!(resolved.selected_root(), &roots[0]); + assert!(Arc::ptr_eq(resolved.environment(), &selected)); + assert!(!Arc::ptr_eq(resolved.environment(), &primary)); + } + #[tokio::test] async fn blocking_snapshot_waits_for_starting_environment() { let listener = TcpListener::bind("127.0.0.1:0") diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index be12e61e08..e006bd8c34 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2828,9 +2828,12 @@ impl Session { .await; } let loaded_agents_md = self.services.agents_md_manager.get_loaded().await; + let selected_capability_roots = environments + .resolve_selected_capability_roots(&self.services.selected_capability_roots); Arc::new(StepContext::new( turn_context, environments, + selected_capability_roots, loaded_agents_md, )) } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 54b531328a..96dffed888 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1047,6 +1047,7 @@ impl Session { // TODO(jif): extract session to share between sub-agents session_extension_data, thread_extension_data, + selected_capability_roots, mcp_thread_init, supports_openai_form_elicitation: std::sync::atomic::AtomicBool::new( supports_openai_form_elicitation, diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index bf672db6e7..e463015459 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -3,12 +3,15 @@ use std::sync::Arc; use crate::agents_md::LoadedAgentsMd; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::session::turn_context::TurnContext; +use codex_exec_server::ResolvedSelectedCapabilityRoot; /// Request-scoped state that may change between model sampling requests. #[derive(Debug)] pub(crate) struct StepContext { pub(crate) turn: Arc, pub(crate) environments: TurnEnvironmentSnapshot, + /// Capability roots bound to ready environments in this exact step. + pub(crate) selected_capability_roots: Vec, /// The canonical AGENTS.md value observed with this environment snapshot. pub(crate) loaded_agents_md: Option>, } @@ -17,11 +20,13 @@ impl StepContext { pub(crate) fn new( turn: Arc, environments: TurnEnvironmentSnapshot, + selected_capability_roots: Vec, loaded_agents_md: Option>, ) -> Self { Self { turn, environments, + selected_capability_roots, loaded_agents_md, } } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index c83b26d098..525fbc3b99 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -190,6 +190,7 @@ impl StepContext { Arc::new(Self::new( turn, environments, + Vec::new(), /*loaded_agents_md*/ None, )) } @@ -5405,6 +5406,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { agent_control.session_id().to_string(), ), thread_extension_data: codex_extension_api::ExtensionData::new(thread_id.to_string()), + selected_capability_roots: Vec::new(), mcp_thread_init: codex_extension_api::ExtensionDataInit::default(), supports_openai_form_elicitation: std::sync::atomic::AtomicBool::new(false), agent_control, @@ -7481,6 +7483,7 @@ where agent_control.session_id().to_string(), ), thread_extension_data: codex_extension_api::ExtensionData::new(thread_id.to_string()), + selected_capability_roots: Vec::new(), mcp_thread_init: codex_extension_api::ExtensionDataInit::default(), supports_openai_form_elicitation: std::sync::atomic::AtomicBool::new(false), agent_control, diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index e19e8c2ff4..cedeaab25e 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -10,6 +10,10 @@ impl Session { step_context: &StepContext, ) -> WorldState { let turn_context = step_context.turn.as_ref(); + tracing::trace!( + selected_capability_root_count = step_context.selected_capability_roots.len(), + "building step world state" + ); let environment_subagents = if turn_context.config.include_environment_context { self.services .agent_control diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index cfed7d90d0..3721cb1848 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -33,6 +33,7 @@ use codex_login::AuthManager; use codex_mcp::McpConnectionManager; use codex_models_manager::manager::SharedModelsManager; use codex_otel::SessionTelemetry; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_rollout::state_db::StateDbHandle; use codex_rollout_trace::ThreadTraceContext; use codex_thread_store::LiveThread; @@ -72,6 +73,9 @@ pub(crate) struct SessionServices { pub(crate) session_extension_data: ExtensionData, pub(crate) thread_extension_data: ExtensionData, pub(crate) supports_openai_form_elicitation: AtomicBool, + /// Raw capability selections for this thread. Each model step resolves them against its + /// current executor environments before using them. + pub(crate) selected_capability_roots: Vec, pub(crate) mcp_thread_init: ExtensionDataInit, pub(crate) agent_control: AgentControl, pub(crate) network_proxy: ArcSwapOption, diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index c422cc749e..6fc5adcb11 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -689,6 +689,7 @@ async fn environment_tools_follow_the_step_context() { let step_context = Arc::new(StepContext::new( Arc::new(turn), environments, + Vec::new(), /*loaded_agents_md*/ None, )); diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 5b2bfa7e55..92c7f85edf 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -22,6 +22,7 @@ mod relay_proto; mod remote; mod remote_file_system; mod remote_process; +mod resolved_capability; mod rpc; mod runtime_paths; mod sandboxed_file_system; @@ -143,6 +144,7 @@ pub use protocol::WriteResponse; pub use protocol::WriteStatus; pub use remote::RemoteEnvironmentConfig; pub use remote::run_remote_environment; +pub use resolved_capability::ResolvedSelectedCapabilityRoot; pub use runtime_paths::ExecServerRuntimePaths; pub use server::DEFAULT_LISTEN_URL; pub use server::ExecServerListenUrlParseError; diff --git a/codex-rs/exec-server/src/resolved_capability.rs b/codex-rs/exec-server/src/resolved_capability.rs new file mode 100644 index 0000000000..5e72b23c5a --- /dev/null +++ b/codex-rs/exec-server/src/resolved_capability.rs @@ -0,0 +1,47 @@ +use std::fmt; +use std::sync::Arc; + +use codex_protocol::capabilities::SelectedCapabilityRoot; + +use crate::Environment; +use crate::ExecutorFileSystem; + +/// A selected capability root pinned to the exact environment instance that owns it. +/// +/// This value is process-local and must not be persisted. Cloning it keeps the same +/// [`Environment`] alive so every consumer of one model step uses the same executor. +#[derive(Clone)] +pub struct ResolvedSelectedCapabilityRoot { + selected_root: SelectedCapabilityRoot, + environment: Arc, +} + +impl ResolvedSelectedCapabilityRoot { + pub fn new(selected_root: SelectedCapabilityRoot, environment: Arc) -> Self { + Self { + selected_root, + environment, + } + } + + pub fn selected_root(&self) -> &SelectedCapabilityRoot { + &self.selected_root + } + + pub fn environment(&self) -> &Arc { + &self.environment + } + + pub fn file_system(&self) -> Arc { + self.environment.get_filesystem() + } +} + +impl fmt::Debug for ResolvedSelectedCapabilityRoot { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ResolvedSelectedCapabilityRoot") + .field("selected_root", &self.selected_root) + .finish_non_exhaustive() + } +}