diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index bc81e9a8a6..b5f442e429 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -27,6 +27,7 @@ pub use codex_config::types::TuiPetAnchor; pub use codex_config::types::UriBasedFileOpener; pub use codex_core::CodexAppsToolsCache; pub use codex_core::CodexThread; +pub use codex_core::EnvironmentConfig; pub use codex_core::ForkSnapshot; pub use codex_core::LoadedAgentsMd; pub use codex_core::McpManager; @@ -88,6 +89,8 @@ pub use codex_model_provider_info::built_in_model_providers; pub use codex_models_manager::manager::RefreshStrategy; pub use codex_models_manager::manager::SharedModelsManager; pub use codex_protocol::ThreadId; +pub use codex_protocol::capabilities::CapabilityRootLocation; +pub use codex_protocol::capabilities::SelectedCapabilityRoot; pub use codex_protocol::config_types::AltScreenMode; pub use codex_protocol::config_types::ApprovalsReviewer; pub use codex_protocol::config_types::AutoCompactTokenLimitScope; diff --git a/codex-rs/core/src/agents_md_tests.rs b/codex-rs/core/src/agents_md_tests.rs index 2c86c3384e..9deefd61e4 100644 --- a/codex-rs/core/src/agents_md_tests.rs +++ b/codex-rs/core/src/agents_md_tests.rs @@ -345,6 +345,7 @@ fn resolved_local_environments( permission_profile: PermissionProfileSnapshot::legacy( PermissionProfile::read_only(), ), + selected_capability_roots: None, }, )) }) diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index b2df5cde07..0c3db6c299 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -1,6 +1,7 @@ use crate::agent::AgentStatus; use crate::config::ConstraintResult; use crate::elicitation::ElicitationRegistration; +use crate::environment_config::EnvironmentConfig; use crate::session::SessionIo; use crate::session::SessionSettingsUpdate; use crate::session::SteerInputError; @@ -766,13 +767,18 @@ impl CodexThread { self.session.thread_environment_selections().await } + /// Installs resolved environment configuration and capability roots on this thread. + pub async fn environment_ready( + &self, + selection: &TurnEnvironmentSelection, + config: EnvironmentConfig, + ) -> CodexResult<()> { + self.session.environment_ready(selection, config).await + } + /// Passively inspects the selected capability roots whose environments are ready now. pub fn inspect_selected_capability_roots(&self) -> SelectedCapabilityRootsStatus { - self.session - .services - .turn_environments - .environment_manager() - .inspect_selected_capability_roots(&self.session.services.selected_capability_roots) + self.session.inspect_selected_capability_roots() } pub async fn read_mcp_resource( diff --git a/codex-rs/core/src/environment_config.rs b/codex-rs/core/src/environment_config.rs new file mode 100644 index 0000000000..6febe40c3b --- /dev/null +++ b/codex-rs/core/src/environment_config.rs @@ -0,0 +1,10 @@ +use codex_protocol::capabilities::SelectedCapabilityRoot; + +/// Resolved configuration supplied by the owner of a thread/environment attachment. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct EnvironmentConfig { + /// Whether shell tools may start login shells in this environment. + pub allow_login_shell: bool, + /// Capability roots selected for this thread's environment attachment. + pub selected_capability_roots: Vec, +} diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 268d0a9261..c3ab8517bb 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -11,6 +11,11 @@ use codex_exec_server::EnvironmentConnectionState; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecServerError; use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::SelectedCapabilityRootsStatus; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; use codex_protocol::protocol::EnvironmentConnectionEvent; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; @@ -22,6 +27,7 @@ use futures::future::BoxFuture; use futures::future::Shared; use tokio_util::task::AbortOnDropHandle; +use crate::environment_config::EnvironmentConfig; use crate::session::turn_context::ShellSnapshotTask; use crate::session::turn_context::TurnEnvironment; use crate::session::turn_context::TurnEnvironmentConfig; @@ -59,7 +65,6 @@ struct ResolvedEnvironment { #[derive(Clone)] struct SelectedTurnEnvironment { selection: TurnEnvironmentSelection, - // Temporary: copied from thread settings until environments supply their own config. config: TurnEnvironmentConfig, environment: Arc, // Selection clones share one listener; the final handle drop aborts it. @@ -74,6 +79,21 @@ pub(crate) struct StartingTurnEnvironment { resolution: TurnEnvironmentResolution, } +impl SelectedTurnEnvironment { + fn has_installed_environment_config(&self) -> bool { + self.config.selected_capability_roots.is_some() + } + + // Thread settings still own permissions, but cannot overwrite shell policy + // that was installed specifically for this environment attachment. + fn update_thread_config(&mut self, config: &TurnEnvironmentConfig) { + if !self.has_installed_environment_config() { + self.config.allow_login_shell = config.allow_login_shell; + } + self.config.permission_profile = config.permission_profile.clone(); + } +} + impl fmt::Debug for StartingTurnEnvironment { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter @@ -118,6 +138,7 @@ impl ThreadEnvironments { }; let selection = environment.selection(); let selected_environment = Arc::clone(&environment.environment); + let inherited_config = environment.config; let resolution: TurnEnvironmentResolution = futures::future::ready(Ok(ResolvedEnvironment { environment: environment.environment, @@ -126,14 +147,17 @@ impl ThreadEnvironments { })) .boxed() .shared(); - Some(SelectedTurnEnvironment { + let mut inherited_environment = SelectedTurnEnvironment { selection, - // Child threads can have different environment config from their parent. - config: thread_environment_config.clone(), + config: inherited_config, environment: selected_environment, connection_events_task: None, resolution, - }) + }; + // Child threads get their own settings, but inherit any + // environment-owned policy that was already installed. + inherited_environment.update_thread_config(&thread_environment_config); + Some(inherited_environment) }) .collect(); Self { @@ -164,7 +188,7 @@ impl ThreadEnvironments { && !matches!(environment.resolution.clone().now_or_never(), Some(Err(_))) { let mut environment = environment.clone(); - environment.config = thread_environment_config.clone(); + environment.update_thread_config(thread_environment_config); next.push(environment); continue; } @@ -234,13 +258,68 @@ impl ThreadEnvironments { .iter() .map(|environment| { let mut environment = environment.clone(); - environment.config = config.clone(); + environment.update_thread_config(config); environment }) .collect(); self.environments.store(Arc::new(environments)); } + /// Installs owner-provided config and roots on their exact thread attachment. + /// Additional environment-owned settings should be applied to its config here. + pub(crate) fn environment_ready( + &self, + selection: &TurnEnvironmentSelection, + config: EnvironmentConfig, + ) -> CodexResult<()> { + let mut environments = Vec::clone(&self.environments.load()); + let Some(environment) = environments.iter_mut().find(|environment| { + environment.selection.environment_id == selection.environment_id + && environment.selection.cwd == selection.cwd + }) else { + return Err(CodexErr::InvalidRequest(format!( + "environment `{}` is not selected on this thread with the requested cwd", + selection.environment_id + ))); + }; + + environment.config.allow_login_shell = config.allow_login_shell; + environment.config.selected_capability_roots = Some(config.selected_capability_roots); + self.environments.store(Arc::new(environments)); + Ok(()) + } + + /// Combines persisted thread roots with installed attachment roots, keeping + /// thread roots first and hiding attachments that are not ready yet. + pub(crate) fn inspect_selected_capability_roots( + &self, + thread_selected_capability_roots: &[SelectedCapabilityRoot], + ) -> SelectedCapabilityRootsStatus { + let environments = self.environments.load(); + let mut selected_capability_roots = thread_selected_capability_roots.to_vec(); + for environment in environments.iter() { + if let Some(roots) = &environment.config.selected_capability_roots { + selected_capability_roots.extend(roots.iter().cloned()); + } + } + let mut seen_root_ids = HashSet::with_capacity(selected_capability_roots.len()); + selected_capability_roots.retain(|root| seen_root_ids.insert(root.id.clone())); + + let mut status = self + .environment_manager + .inspect_selected_capability_roots(&selected_capability_roots); + status.ready_roots.retain(|root| { + let CapabilityRootLocation::Environment { environment_id, .. } = &root.location; + environments + .iter() + .find(|environment| &environment.selection.environment_id == environment_id) + .is_none_or(|environment| { + matches!(environment.resolution.clone().now_or_never(), Some(Ok(_))) + }) + }); + status + } + fn spawn_connection_event_listener( environment: &Environment, environment_id: String, @@ -558,6 +637,7 @@ mod tests { TurnEnvironmentConfig { allow_login_shell: true, permission_profile: PermissionProfileSnapshot::legacy(PermissionProfile::read_only()), + selected_capability_roots: None, } } @@ -725,6 +805,7 @@ url = "ws://127.0.0.1:8765" ActivePermissionProfile::read_only(), vec![cwd.join("profile-root")], ), + selected_capability_roots: None, }; let turn_environments = ThreadEnvironments::new( Arc::new(EnvironmentManager::default_for_tests()), @@ -917,6 +998,7 @@ url = "ws://127.0.0.1:8765" ActivePermissionProfile::read_only(), vec![cwd.join("profile-root")], ), + selected_capability_roots: None, }; let cwd = PathUri::from_abs_path(&cwd); let remote = TurnEnvironmentSelection { @@ -1042,7 +1124,28 @@ url = "ws://127.0.0.1:8765" environments .update_selections(std::slice::from_ref(&selection), &test_environment_config()); let failed_resolution = environments.environments.load()[0].resolution.clone(); - assert!(failed_resolution.clone().await.is_err()); + let error = failed_resolution + .clone() + .await + .err() + .expect("environment should fail to start"); + let selected_root = SelectedCapabilityRoot { + id: "failed-root".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: selection.environment_id.clone(), + path: selection.cwd.clone(), + }, + }; + assert_eq!( + environments.inspect_selected_capability_roots(&[selected_root]), + SelectedCapabilityRootsStatus { + ready_roots: Vec::new(), + warnings: vec![format!( + "selected capability environment `{}` is unavailable: {error}", + selection.environment_id + )], + } + ); let listener = TcpListener::bind("127.0.0.1:0") .await @@ -1211,6 +1314,7 @@ url = "ws://127.0.0.1:8765" ActivePermissionProfile::read_only(), vec![cwd.join("child-profile-root")], ), + selected_capability_roots: None, }; let environments = ThreadEnvironments::new( manager, @@ -1231,6 +1335,92 @@ url = "ws://127.0.0.1:8765" assert_eq!(inherited.config, child_config); } + #[tokio::test] + async fn installed_environment_config_is_inherited_and_reset_for_new_cwd() { + let cwd = AbsolutePathBuf::current_dir().expect("cwd"); + let selection = TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: PathUri::from_abs_path(&cwd), + workspace_roots: Vec::new(), + }; + let root = |id: &str| SelectedCapabilityRoot { + id: id.to_string(), + location: CapabilityRootLocation::Environment { + environment_id: selection.environment_id.clone(), + path: selection.cwd.clone(), + }, + }; + let manager = Arc::new(EnvironmentManager::default_for_tests()); + let parent = + resolve_turn_environments(Arc::clone(&manager), std::slice::from_ref(&selection)).await; + parent + .environment_ready( + &selection, + EnvironmentConfig { + allow_login_shell: false, + selected_capability_roots: vec![root("parent-root")], + }, + ) + .expect("install environment config"); + + let child_thread_config = TurnEnvironmentConfig { + permission_profile: PermissionProfileSnapshot::legacy( + PermissionProfile::workspace_write(), + ), + ..test_environment_config() + }; + let child = ThreadEnvironments::new( + manager, + crate::shell::default_user_shell(), + child_thread_config.clone(), + ShellSnapshot::disabled(), + parent.snapshot().await, + /*non_blocking_snapshots*/ false, + ); + let child_snapshot = child.snapshot().await; + let child_environment = child_snapshot.primary().expect("child environment"); + + parent + .environment_ready( + &selection, + EnvironmentConfig { + allow_login_shell: false, + selected_capability_roots: Vec::new(), + }, + ) + .expect("clear parent roots"); + let cleared_snapshot = parent.snapshot().await; + assert_eq!( + cleared_snapshot + .primary() + .expect("environment with cleared roots") + .config, + TurnEnvironmentConfig { + allow_login_shell: false, + selected_capability_roots: Some(Vec::new()), + ..test_environment_config() + } + ); + assert_eq!( + child_environment.config, + TurnEnvironmentConfig { + allow_login_shell: false, + selected_capability_roots: Some(vec![root("parent-root")]), + ..child_thread_config + } + ); + + let changed_selection = TurnEnvironmentSelection { + cwd: PathUri::from_abs_path(&cwd.join("changed")), + ..selection + }; + let new_thread_config = test_environment_config(); + parent.update_selections(std::slice::from_ref(&changed_selection), &new_thread_config); + let changed_snapshot = parent.snapshot().await; + let changed_environment = changed_snapshot.primary().expect("changed environment"); + assert_eq!(changed_environment.config, new_thread_config); + } + #[tokio::test] async fn single_local_environment_cwd_requires_exactly_one_local_environment() { let cwd = AbsolutePathBuf::current_dir().expect("cwd"); diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index d8db808548..b4aebe4a3c 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -32,6 +32,7 @@ pub use codex_thread::CodexThreadSettingsOverrides; pub use codex_thread::ThreadConfigSnapshot; pub use codex_thread::TryStartTurnIfIdleError; pub use codex_thread::TryStartTurnIfIdleRejectionReason; +pub use environment_config::EnvironmentConfig; pub use session::turn_context::TurnContext; pub use user_message_admission::UserMessageAdmission; pub use user_message_admission::UserMessageAdmissionError; @@ -46,6 +47,7 @@ pub mod context; mod context_manager; mod current_time; mod elicitation; +mod environment_config; mod environment_selection; pub mod exec; pub mod exec_env; diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 74a8f17ba6..7fa18ef89c 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1107,6 +1107,7 @@ async fn mcp_sandbox_cwd_uses_matching_server_environment_uri() -> anyhow::Resul .permissions .permission_profile_state() .snapshot(), + selected_capability_roots: None, }, ))); diff --git a/codex-rs/core/src/session/environment.rs b/codex-rs/core/src/session/environment.rs new file mode 100644 index 0000000000..fe9fa0504b --- /dev/null +++ b/codex-rs/core/src/session/environment.rs @@ -0,0 +1,55 @@ +use std::collections::HashSet; + +use codex_exec_server::MAX_SELECTED_CAPABILITY_ROOTS; +use codex_exec_server::SelectedCapabilityRootsStatus; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; +use codex_protocol::protocol::TurnEnvironmentSelection; + +use crate::environment_config::EnvironmentConfig; +use crate::session::session::Session; + +impl Session { + pub(crate) async fn environment_ready( + &self, + selection: &TurnEnvironmentSelection, + config: EnvironmentConfig, + ) -> CodexResult<()> { + if config.selected_capability_roots.len() > MAX_SELECTED_CAPABILITY_ROOTS { + return Err(CodexErr::InvalidRequest(format!( + "environment readiness contains more than {MAX_SELECTED_CAPABILITY_ROOTS} selected capability roots" + ))); + } + + let mut root_ids = HashSet::with_capacity(config.selected_capability_roots.len()); + for root in &config.selected_capability_roots { + let CapabilityRootLocation::Environment { environment_id, .. } = &root.location; + if root.id.trim().is_empty() + || environment_id != &selection.environment_id + || !root_ids.insert(root.id.as_str()) + { + return Err(CodexErr::InvalidRequest(format!( + "selected capability roots must have unique non-empty IDs and belong to environment `{}`", + selection.environment_id + ))); + } + } + + // grab session lock so installation can't race w/ thread settings updates + let _state = self.state.lock().await; + self.services + .turn_environments + .environment_ready(selection, config)?; + // mark mcp runtime for refresh because available capabilities could've changed + self.mark_mcp_runtime_dirty(); + Ok(()) + } + + /// Combines this session's persisted roots with ready environment attachments. + pub(crate) fn inspect_selected_capability_roots(&self) -> SelectedCapabilityRootsStatus { + self.services + .turn_environments + .inspect_selected_capability_roots(&self.services.selected_capability_roots) + } +} diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 77f95df876..d3de7ca2f8 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -442,11 +442,13 @@ impl Session { .selected_capability_roots .iter() .cloned() - .chain( - environments - .turn_environments() - .flat_map(|environment| environment.environment.selected_capability_roots()), - ) + .chain(environments.turn_environments().flat_map(|environment| { + environment + .config + .selected_capability_roots + .clone() + .unwrap_or_else(|| environment.environment.selected_capability_roots()) + })) .enumerate() { if let Some(kept_location) = root_locations_by_id.get(&root.id) { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index b2e5a3987e..9961b0b53d 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -205,6 +205,7 @@ use codex_protocol::exec_output::StreamOutput; mod code_mode_warning; pub(crate) mod context_window; +mod environment; mod extension_metrics; mod handlers; mod inject; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 3b75e06bd9..bfe89311fe 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -164,6 +164,7 @@ impl SessionConfiguration { .permissions .allow_login_shell, permission_profile: self.permission_profile_state.snapshot(), + selected_capability_roots: None, } } diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index f79840d617..60f1d14cc3 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -10,6 +10,7 @@ use codex_file_system::FileSystemSandboxContext; use codex_model_provider::SharedModelProvider; use codex_protocol::SessionId; use codex_protocol::ThreadId; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::openai_models::MODEL_SPECIALTY_CYBER; use codex_protocol::openai_models::ModelInfo; @@ -34,6 +35,8 @@ pub(crate) type ShellSnapshotTask = Shared>, } #[derive(Clone)] diff --git a/codex-rs/core/src/tools/handlers/shell_tests.rs b/codex-rs/core/src/tools/handlers/shell_tests.rs index 5f8559a143..7efa45c741 100644 --- a/codex-rs/core/src/tools/handlers/shell_tests.rs +++ b/codex-rs/core/src/tools/handlers/shell_tests.rs @@ -120,6 +120,7 @@ async fn shell_command_handler_to_exec_params_uses_selected_environment() { permission_profile, active_permission_profile.clone(), ), + selected_capability_roots: None, }, ); let mut expected_env = create_env( diff --git a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs index 8be5183e67..fa9fbc0304 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs @@ -26,6 +26,7 @@ fn test_turn_environment(environment_id: &str) -> crate::session::turn_context:: TurnEnvironmentConfig { allow_login_shell: true, permission_profile: PermissionProfileSnapshot::legacy(PermissionProfile::read_only()), + selected_capability_roots: None, }, ) } diff --git a/codex-rs/core/src/tools/runtimes/shell_tests.rs b/codex-rs/core/src/tools/runtimes/shell_tests.rs index ac36ecd280..6b923af88d 100644 --- a/codex-rs/core/src/tools/runtimes/shell_tests.rs +++ b/codex-rs/core/src/tools/runtimes/shell_tests.rs @@ -25,6 +25,7 @@ async fn approval_key_uses_path_uri_and_includes_environment_id() { permission_profile: PermissionProfileSnapshot::legacy( PermissionProfile::read_only(), ), + selected_capability_roots: None, }, ), shell_type: None, diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index ae4cb09d06..382ebc7b6e 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -507,6 +507,7 @@ mod tests { permission_profile: PermissionProfileSnapshot::legacy( PermissionProfile::read_only(), ), + selected_capability_roots: None, }, ) } diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 07874e89c5..0d2aef0357 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -919,6 +919,7 @@ async fn zsh_fork_unified_exec_keeps_shell_parameter_when_remote_environment_ava .permissions .permission_profile_state() .snapshot(), + selected_capability_roots: None, }, ), )); diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index b8187490dc..3702838d37 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -4,6 +4,8 @@ use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_api::AuthProvider; use codex_config::types::ApprovalsReviewer; +use codex_core::EnvironmentConfig; +use codex_core::StartThreadOptions; use codex_core::WaitForEnvironmentToolConfig; use codex_core::compact::SUMMARIZATION_PROMPT; use codex_core::config::Config; @@ -22,6 +24,7 @@ use codex_exec_server::REMOTE_ENVIRONMENT_ID; use codex_exec_server::RemoteEnvironmentConfig; use codex_exec_server::RemoveOptions; use codex_extension_api::ContextContributor; +use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::RenderedWorldStateFragment; @@ -953,6 +956,164 @@ async fn wait_for_response_request_count(response_mock: &ResponseMock, expected_ .expect("timed out waiting for Responses API request"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shared_executor_keeps_ready_capability_roots_scoped_to_each_attachment() -> Result<()> { + let server = start_mock_server().await; + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension)); + let mut builder = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| { + config.use_experimental_unified_exec_tool = true; + assert!(config.features.enable(Feature::UnifiedExec).is_ok()); + }); + let test = builder.build_with_auto_env(&server).await?; + let selection = test + .codex + .environment_selections() + .await + .into_iter() + .next() + .context("thread should select its executor environment")?; + let root = |id: &str| SelectedCapabilityRoot { + id: id.to_string(), + location: CapabilityRootLocation::Environment { + environment_id: selection.environment_id.clone(), + path: selection.cwd.clone(), + }, + }; + + let mut second_thread_init = ExtensionDataInit::new(); + second_thread_init.insert(vec![root("startup-root")]); + let second = test + .thread_manager + .start_thread(StartThreadOptions { + environments: Some(vec![selection.clone()]), + thread_extension_init: second_thread_init, + ..StartThreadOptions::new(test.config.clone()) + }) + .await?; + + test.codex + .environment_ready( + &selection, + EnvironmentConfig { + allow_login_shell: false, + selected_capability_roots: vec![root("first-root")], + }, + ) + .await?; + second + .thread + .environment_ready( + &selection, + EnvironmentConfig { + allow_login_shell: true, + selected_capability_roots: vec![root("startup-root"), root("second-root")], + }, + ) + .await?; + + assert_eq!( + test.codex.inspect_selected_capability_roots().ready_roots, + vec![root("first-root")] + ); + assert_eq!( + second + .thread + .inspect_selected_capability_roots() + .ready_roots, + vec![root("startup-root"), root("second-root")] + ); + + let response_mock = mount_sse_sequence( + &server, + ["first", "second", "first-updated", "second-again"] + .into_iter() + .map(|response_id| { + sse(vec![ + ev_response_created(response_id), + ev_completed(response_id), + ]) + }) + .collect(), + ) + .await; + + for (index, (thread, prompt)) in [ + (&test.codex, "first"), + (&second.thread, "second"), + (&test.codex, "first-updated"), + (&second.thread, "second-again"), + ] + .into_iter() + .enumerate() + { + if index == 2 { + test.codex + .environment_ready( + &selection, + EnvironmentConfig { + allow_login_shell: false, + selected_capability_roots: vec![root("first-updated-root")], + }, + ) + .await?; + } + + thread + .submit( + vec![UserInput::Text { + text: prompt.to_string(), + text_elements: Vec::new(), + }] + .into(), + ) + .await?; + wait_for_event(thread, |event| matches!(event, EventMsg::TurnComplete(_))).await; + } + + let requests = response_mock.requests(); + let root_fragments = requests + .iter() + .map(|request| { + request + .message_input_texts("user") + .into_iter() + .rfind(|text| text.contains("")) + .context("ready capability roots should be model-visible") + }) + .collect::>>()?; + assert_eq!( + root_fragments, + vec![ + "first-root", + "startup-root,second-root", + "first-updated-root", + "startup-root,second-root", + ] + ); + + let login_shells = requests + .iter() + .map(|request| { + let body = request.body_json(); + let exec_command = body["tools"] + .as_array() + .context("tools should be an array")? + .iter() + .find(|tool| tool["name"] == "exec_command") + .context("exec_command should be available")?; + Ok(exec_command["parameters"]["properties"] + .get("login") + .is_some()) + }) + .collect::>>()?; + assert_eq!(login_shells, vec![false, true, false, true]); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn ready_before_selection_exposes_remote_tools_and_capability_context_after_wait() -> Result<()> {