Bind selected capability roots to model steps

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<ResolvedSelectedCapabilityRoot> {
+        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")

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,
         ))
     }

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,

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<TurnContext>,
     pub(crate) environments: TurnEnvironmentSnapshot,
+    /// Capability roots bound to ready environments in this exact step.
+    pub(crate) selected_capability_roots: Vec<ResolvedSelectedCapabilityRoot>,
     /// The canonical AGENTS.md value observed with this environment snapshot.
     pub(crate) loaded_agents_md: Option<Arc<LoadedAgentsMd>>,
 }
@@ -17,11 +20,13 @@ impl StepContext {
     pub(crate) fn new(
         turn: Arc<TurnContext>,
         environments: TurnEnvironmentSnapshot,
+        selected_capability_roots: Vec<ResolvedSelectedCapabilityRoot>,
         loaded_agents_md: Option<Arc<LoadedAgentsMd>>,
     ) -> Self {
         Self {
             turn,
             environments,
+            selected_capability_roots,
             loaded_agents_md,
         }
     }

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,

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

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<SelectedCapabilityRoot>,
     pub(crate) mcp_thread_init: ExtensionDataInit,
     pub(crate) agent_control: AgentControl,
     pub(crate) network_proxy: ArcSwapOption<StartedNetworkProxy>,

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,
     ));
 

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;

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<Environment>,
+}
+
+impl ResolvedSelectedCapabilityRoot {
+    pub fn new(selected_root: SelectedCapabilityRoot, environment: Arc<Environment>) -> Self {
+        Self {
+            selected_root,
+            environment,
+        }
+    }
+
+    pub fn selected_root(&self) -> &SelectedCapabilityRoot {
+        &self.selected_root
+    }
+
+    pub fn environment(&self) -> &Arc<Environment> {
+        &self.environment
+    }
+
+    pub fn file_system(&self) -> Arc<dyn ExecutorFileSystem> {
+        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()
+    }
+}