From 03d55f0e6f5a5006270340df7827a91972b1caa4 Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Wed, 4 Mar 2026 23:02:27 -0800 Subject: [PATCH 1/4] chore: add web_search_tool_type for image support (#13538) add `web_search_tool_type` on model_info that can be populated from backend. will be used to filter which models can use `web_search` with images and which cant. added small unit test. --- .../app-server/tests/common/models_cache.rs | 1 + .../codex-api/tests/models_integration.rs | 1 + codex-rs/core/src/client_common.rs | 2 + .../core/src/models_manager/model_info.rs | 2 + codex-rs/core/src/tools/spec.rs | 71 +++++++++++++++---- codex-rs/core/tests/suite/model_switching.rs | 2 + codex-rs/core/tests/suite/models_cache_ttl.rs | 1 + codex-rs/core/tests/suite/personality.rs | 2 + codex-rs/core/tests/suite/remote_models.rs | 3 + codex-rs/core/tests/suite/rmcp_client.rs | 1 + codex-rs/core/tests/suite/view_image.rs | 1 + codex-rs/protocol/src/openai_models.rs | 14 ++++ 12 files changed, 89 insertions(+), 12 deletions(-) diff --git a/codex-rs/app-server/tests/common/models_cache.rs b/codex-rs/app-server/tests/common/models_cache.rs index 0c2a2e5857..62a8f55d1b 100644 --- a/codex-rs/app-server/tests/common/models_cache.rs +++ b/codex-rs/app-server/tests/common/models_cache.rs @@ -36,6 +36,7 @@ fn preset_to_info(preset: &ModelPreset, priority: i32) -> ModelInfo { default_verbosity: None, availability_nux: None, apply_patch_tool_type: None, + web_search_tool_type: Default::default(), truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, diff --git a/codex-rs/codex-api/tests/models_integration.rs b/codex-rs/codex-api/tests/models_integration.rs index 3d9133700c..cc471ae890 100644 --- a/codex-rs/codex-api/tests/models_integration.rs +++ b/codex-rs/codex-api/tests/models_integration.rs @@ -84,6 +84,7 @@ async fn models_client_hits_models_endpoint() { default_verbosity: None, availability_nux: None, apply_patch_tool_type: None, + web_search_tool_type: Default::default(), truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index 50facf53e4..d775c9a376 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -176,6 +176,8 @@ pub(crate) mod tools { WebSearch { #[serde(skip_serializing_if = "Option::is_none")] external_web_access: Option, + #[serde(skip_serializing_if = "Option::is_none")] + search_content_types: Option>, }, #[serde(rename = "custom")] Freeform(FreeformTool), diff --git a/codex-rs/core/src/models_manager/model_info.rs b/codex-rs/core/src/models_manager/model_info.rs index 643d42571d..3664a52660 100644 --- a/codex-rs/core/src/models_manager/model_info.rs +++ b/codex-rs/core/src/models_manager/model_info.rs @@ -6,6 +6,7 @@ use codex_protocol::openai_models::ModelMessages; use codex_protocol::openai_models::ModelVisibility; use codex_protocol::openai_models::TruncationMode; use codex_protocol::openai_models::TruncationPolicyConfig; +use codex_protocol::openai_models::WebSearchToolType; use codex_protocol::openai_models::default_input_modalities; use crate::config::Config; @@ -78,6 +79,7 @@ pub(crate) fn model_info_from_slug(slug: &str) -> ModelInfo { support_verbosity: false, default_verbosity: None, apply_patch_tool_type: None, + web_search_tool_type: WebSearchToolType::Text, truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index dcdb48500e..00867b20d8 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -25,6 +25,7 @@ use codex_protocol::openai_models::ApplyPatchToolType; use codex_protocol::openai_models::ConfigShellToolType; use codex_protocol::openai_models::InputModality; use codex_protocol::openai_models::ModelInfo; +use codex_protocol::openai_models::WebSearchToolType; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use serde::Deserialize; @@ -36,6 +37,7 @@ use std::collections::HashMap; const SEARCH_TOOL_BM25_DESCRIPTION_TEMPLATE: &str = include_str!("../../templates/search_tool/tool_description.md"); +const WEB_SEARCH_CONTENT_TYPES: [&str; 2] = ["text", "image"]; #[derive(Debug, Clone, Copy, Eq, PartialEq)] pub enum ShellCommandBackendConfig { Classic, @@ -49,6 +51,7 @@ pub(crate) struct ToolsConfig { pub allow_login_shell: bool, pub apply_patch_tool_type: Option, pub web_search_mode: Option, + pub web_search_tool_type: WebSearchToolType, pub image_gen_tool: bool, pub agent_roles: BTreeMap, pub search_tool: bool, @@ -140,6 +143,7 @@ impl ToolsConfig { allow_login_shell: true, apply_patch_tool_type, web_search_mode: *web_search_mode, + web_search_tool_type: model_info.web_search_tool_type, image_gen_tool: include_image_gen_tool, agent_roles: BTreeMap::new(), search_tool: include_search_tool, @@ -1877,18 +1881,27 @@ pub(crate) fn build_specs( builder.register_handler("test_sync_tool", test_sync_handler); } - match config.web_search_mode { - Some(WebSearchMode::Cached) => { - builder.push_spec(ToolSpec::WebSearch { - external_web_access: Some(false), - }); - } - Some(WebSearchMode::Live) => { - builder.push_spec(ToolSpec::WebSearch { - external_web_access: Some(true), - }); - } - Some(WebSearchMode::Disabled) | None => {} + let external_web_access = match config.web_search_mode { + Some(WebSearchMode::Cached) => Some(false), + Some(WebSearchMode::Live) => Some(true), + Some(WebSearchMode::Disabled) | None => None, + }; + + if let Some(external_web_access) = external_web_access { + let search_content_types = match config.web_search_tool_type { + WebSearchToolType::Text => None, + WebSearchToolType::TextAndImage => Some( + WEB_SEARCH_CONTENT_TYPES + .into_iter() + .map(str::to_string) + .collect(), + ), + }; + + builder.push_spec(ToolSpec::WebSearch { + external_web_access: Some(external_web_access), + search_content_types, + }); } if config.image_gen_tool { @@ -2172,6 +2185,7 @@ mod tests { create_apply_patch_freeform_tool(), ToolSpec::WebSearch { external_web_access: Some(true), + search_content_types: None, }, create_view_image_tool(), ] { @@ -2486,6 +2500,7 @@ mod tests { tool.spec, ToolSpec::WebSearch { external_web_access: Some(false), + search_content_types: None, } ); } @@ -2510,6 +2525,38 @@ mod tests { tool.spec, ToolSpec::WebSearch { external_web_access: Some(true), + search_content_types: None, + } + ); + } + + #[test] + fn web_search_tool_type_text_and_image_sets_search_content_types() { + let config = test_config(); + let mut model_info = + ModelsManager::construct_model_info_offline_for_tests("gpt-5-codex", &config); + model_info.web_search_tool_type = WebSearchToolType::TextAndImage; + let features = Features::with_defaults(); + + let tools_config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + features: &features, + web_search_mode: Some(WebSearchMode::Live), + session_source: SessionSource::Cli, + }); + let (tools, _) = build_specs(&tools_config, None, None, &[]).build(); + + let tool = find_tool(&tools, "web_search"); + assert_eq!( + tool.spec, + ToolSpec::WebSearch { + external_web_access: Some(true), + search_content_types: Some( + WEB_SEARCH_CONTENT_TYPES + .into_iter() + .map(str::to_string) + .collect() + ), } ); } diff --git a/codex-rs/core/tests/suite/model_switching.rs b/codex-rs/core/tests/suite/model_switching.rs index a93d0f751b..937de8d536 100644 --- a/codex-rs/core/tests/suite/model_switching.rs +++ b/codex-rs/core/tests/suite/model_switching.rs @@ -65,6 +65,7 @@ fn test_model_info( default_verbosity: None, availability_nux: None, apply_patch_tool_type: None, + web_search_tool_type: Default::default(), truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, @@ -684,6 +685,7 @@ async fn model_switch_to_smaller_model_updates_token_context_window() -> Result< default_verbosity: None, availability_nux: None, apply_patch_tool_type: None, + web_search_tool_type: Default::default(), truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, diff --git a/codex-rs/core/tests/suite/models_cache_ttl.rs b/codex-rs/core/tests/suite/models_cache_ttl.rs index 54b0629243..e8f9cbf7fe 100644 --- a/codex-rs/core/tests/suite/models_cache_ttl.rs +++ b/codex-rs/core/tests/suite/models_cache_ttl.rs @@ -342,6 +342,7 @@ fn test_remote_model(slug: &str, priority: i32) -> ModelInfo { default_verbosity: None, availability_nux: None, apply_patch_tool_type: None, + web_search_tool_type: Default::default(), truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, diff --git a/codex-rs/core/tests/suite/personality.rs b/codex-rs/core/tests/suite/personality.rs index 548c09ba45..754c46ebfb 100644 --- a/codex-rs/core/tests/suite/personality.rs +++ b/codex-rs/core/tests/suite/personality.rs @@ -647,6 +647,7 @@ async fn remote_model_friendly_personality_instructions_with_feature() -> anyhow default_verbosity: None, availability_nux: None, apply_patch_tool_type: None, + web_search_tool_type: Default::default(), truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, @@ -761,6 +762,7 @@ async fn user_turn_personality_remote_model_template_includes_update_message() - default_verbosity: None, availability_nux: None, apply_patch_tool_type: None, + web_search_tool_type: Default::default(), truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, diff --git a/codex-rs/core/tests/suite/remote_models.rs b/codex-rs/core/tests/suite/remote_models.rs index ac2bc71565..4610bec096 100644 --- a/codex-rs/core/tests/suite/remote_models.rs +++ b/codex-rs/core/tests/suite/remote_models.rs @@ -301,6 +301,7 @@ async fn remote_models_remote_model_uses_unified_exec() -> Result<()> { default_verbosity: None, availability_nux: None, apply_patch_tool_type: None, + web_search_tool_type: Default::default(), truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, @@ -542,6 +543,7 @@ async fn remote_models_apply_remote_base_instructions() -> Result<()> { default_verbosity: None, availability_nux: None, apply_patch_tool_type: None, + web_search_tool_type: Default::default(), truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, @@ -1007,6 +1009,7 @@ fn test_remote_model_with_policy( default_verbosity: None, availability_nux: None, apply_patch_tool_type: None, + web_search_tool_type: Default::default(), truncation_policy, supports_parallel_tool_calls: false, supports_image_detail_original: false, diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index a8d414e998..4139a582a8 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -407,6 +407,7 @@ async fn stdio_image_responses_are_sanitized_for_text_only_model() -> anyhow::Re default_verbosity: None, availability_nux: None, apply_patch_tool_type: None, + web_search_tool_type: Default::default(), truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, diff --git a/codex-rs/core/tests/suite/view_image.rs b/codex-rs/core/tests/suite/view_image.rs index 3e652577bb..3a3c89ac75 100644 --- a/codex-rs/core/tests/suite/view_image.rs +++ b/codex-rs/core/tests/suite/view_image.rs @@ -1003,6 +1003,7 @@ async fn view_image_tool_returns_unsupported_message_for_text_only_model() -> an default_verbosity: None, availability_nux: None, apply_patch_tool_type: None, + web_search_tool_type: Default::default(), truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, diff --git a/codex-rs/protocol/src/openai_models.rs b/codex-rs/protocol/src/openai_models.rs index 5968ae699b..3d668c4477 100644 --- a/codex-rs/protocol/src/openai_models.rs +++ b/codex-rs/protocol/src/openai_models.rs @@ -181,6 +181,16 @@ pub enum ApplyPatchToolType { Function, } +#[derive( + Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash, TS, JsonSchema, Default, +)] +#[serde(rename_all = "snake_case")] +pub enum WebSearchToolType { + #[default] + Text, + TextAndImage, +} + /// Server-provided truncation policy metadata for a model. #[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, TS, JsonSchema)] #[serde(rename_all = "snake_case")] @@ -243,6 +253,8 @@ pub struct ModelInfo { pub support_verbosity: bool, pub default_verbosity: Option, pub apply_patch_tool_type: Option, + #[serde(default)] + pub web_search_tool_type: WebSearchToolType, pub truncation_policy: TruncationPolicyConfig, pub supports_parallel_tool_calls: bool, #[serde(default)] @@ -515,6 +527,7 @@ mod tests { support_verbosity: false, default_verbosity: None, apply_patch_tool_type: None, + web_search_tool_type: WebSearchToolType::Text, truncation_policy: TruncationPolicyConfig::bytes(10_000), supports_parallel_tool_calls: false, supports_image_detail_original: false, @@ -718,6 +731,7 @@ mod tests { assert_eq!(model.availability_nux, None); assert!(!model.supports_image_detail_original); + assert_eq!(model.web_search_tool_type, WebSearchToolType::Text); } #[test] From 1ce1712aebb457febfe1bbb96b9b33f8a617cec2 Mon Sep 17 00:00:00 2001 From: pash-openai Date: Thu, 5 Mar 2026 00:00:16 -0800 Subject: [PATCH 2/4] [tui] Show speed in session header (#13446) - add a speed row to the startup/session header under the model row - render the speed row with the same styling pattern as the model row, using /fast to change - show only Fast or Standard to users and update the affected snapshots --------- Co-authored-by: Codex --- .../app-server/src/codex_message_processor.rs | 45 +++++++++++++++++++ codex-rs/tui/src/app.rs | 4 ++ codex-rs/tui/src/chatwidget.rs | 12 +++++ codex-rs/tui/src/chatwidget/tests.rs | 20 +++++++++ codex-rs/tui/src/history_cell.rs | 40 ++++++++++++++++- ...ter_long_transcript_fresh_header_only.snap | 1 - ...nfo_availability_nux_tooltip_snapshot.snap | 1 - 7 files changed, 119 insertions(+), 4 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 2f36602393..44a74a0675 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -6493,6 +6493,14 @@ fn collect_resume_override_mismatches( config_snapshot.model_provider_id )); } + if let Some(requested_service_tier) = request.service_tier.as_ref() + && requested_service_tier != &config_snapshot.service_tier + { + mismatch_details.push(format!( + "service_tier requested={requested_service_tier:?} active={:?}", + config_snapshot.service_tier + )); + } if let Some(requested_cwd) = request.cwd.as_deref() { let requested_cwd_path = std::path::PathBuf::from(requested_cwd); if requested_cwd_path != config_snapshot.cwd { @@ -7259,6 +7267,43 @@ mod tests { validate_dynamic_tools(&tools).expect("valid schema"); } + #[test] + fn collect_resume_override_mismatches_includes_service_tier() { + let request = ThreadResumeParams { + thread_id: "thread-1".to_string(), + history: None, + path: None, + model: None, + model_provider: None, + service_tier: Some(Some(codex_protocol::config_types::ServiceTier::Fast)), + cwd: None, + approval_policy: None, + sandbox: None, + config: None, + base_instructions: None, + developer_instructions: None, + personality: None, + persist_extended_history: false, + }; + let config_snapshot = ThreadConfigSnapshot { + model: "gpt-5".to_string(), + model_provider_id: "openai".to_string(), + service_tier: Some(codex_protocol::config_types::ServiceTier::Flex), + approval_policy: codex_protocol::protocol::AskForApproval::OnRequest, + sandbox_policy: codex_protocol::protocol::SandboxPolicy::DangerFullAccess, + cwd: PathBuf::from("/tmp"), + ephemeral: false, + reasoning_effort: None, + personality: None, + session_source: SessionSource::Cli, + }; + + assert_eq!( + collect_resume_override_mismatches(&request, &config_snapshot), + vec!["service_tier requested=Some(Fast) active=Some(Flex)".to_string()] + ); + } + #[test] fn extract_conversation_summary_prefers_plain_user_messages() -> Result<()> { let conversation_id = ThreadId::from_string("3f941c35-29b3-493b-b0a4-e25800d9aeb0")?; diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index fe8aca54fd..5722c5020a 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -806,6 +806,8 @@ impl App { history_cell::SessionHeaderHistoryCell::new( self.chat_widget.current_model().to_string(), self.chat_widget.current_reasoning_effort(), + self.chat_widget + .should_show_fast_status(self.chat_widget.current_service_tier()), self.config.cwd.clone(), version, ) @@ -5178,6 +5180,7 @@ mod tests { is_first, None, None, + false, )) as Arc }; @@ -5827,6 +5830,7 @@ mod tests { is_first, None, None, + false, )) as Arc }; diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 7a71d40647..5d0a9dbec2 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -1224,6 +1224,7 @@ impl ChatWidget { self.sync_fast_command_enabled(); self.sync_personality_command_enabled(); let startup_tooltip_override = self.startup_tooltip_override.take(); + let show_fast_status = self.should_show_fast_status(event.service_tier); let session_info_cell = history_cell::new_session_info( &self.config, &model_for_header, @@ -1233,6 +1234,7 @@ impl ChatWidget { self.auth_manager .auth_cached() .and_then(|auth| auth.account_plan_type()), + show_fast_status, ); self.apply_session_info_cell(session_info_cell); @@ -7221,6 +7223,15 @@ impl ChatWidget { self.config.service_tier } + pub(crate) fn should_show_fast_status(&self, service_tier: Option) -> bool { + matches!(service_tier, Some(ServiceTier::Fast)) + && self + .auth_manager + .auth_cached() + .as_ref() + .is_some_and(CodexAuth::is_chatgpt_auth) + } + fn fast_mode_enabled(&self) -> bool { self.config.features.enabled(Feature::FastMode) } @@ -7555,6 +7566,7 @@ impl ChatWidget { DEFAULT_MODEL_DISPLAY_NAME.to_string(), placeholder_style, None, + false, config.cwd.clone(), CODEX_CLI_VERSION, )) diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index 72fa61c13a..6d8ecc78f1 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -7458,6 +7458,26 @@ async fn user_turn_carries_service_tier_after_fast_toggle() { } } +#[tokio::test] +async fn fast_status_indicator_requires_chatgpt_auth() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5.3-codex")).await; + chat.set_service_tier(Some(ServiceTier::Fast)); + + assert!(!chat.should_show_fast_status(chat.current_service_tier())); + + set_chatgpt_auth(&mut chat); + + assert!(chat.should_show_fast_status(chat.current_service_tier())); +} + +#[tokio::test] +async fn fast_status_indicator_is_hidden_when_fast_mode_is_off() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5.3-codex")).await; + set_chatgpt_auth(&mut chat); + + assert!(!chat.should_show_fast_status(chat.current_service_tier())); +} + #[tokio::test] async fn approvals_popup_shows_disabled_presets() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(None).await; diff --git a/codex-rs/tui/src/history_cell.rs b/codex-rs/tui/src/history_cell.rs index b5e839c32f..3a568a0f8f 100644 --- a/codex-rs/tui/src/history_cell.rs +++ b/codex-rs/tui/src/history_cell.rs @@ -1047,6 +1047,7 @@ pub(crate) fn new_session_info( is_first_event: bool, tooltip_override: Option, auth_plan: Option, + show_fast_status: bool, ) -> SessionInfoCell { let SessionConfiguredEvent { model, @@ -1057,6 +1058,7 @@ pub(crate) fn new_session_info( let header = SessionHeaderHistoryCell::new( model.clone(), reasoning_effort, + show_fast_status, config.cwd.clone(), CODEX_CLI_VERSION, ); @@ -1143,6 +1145,7 @@ pub(crate) struct SessionHeaderHistoryCell { model: String, model_style: Style, reasoning_effort: Option, + show_fast_status: bool, directory: PathBuf, } @@ -1150,6 +1153,7 @@ impl SessionHeaderHistoryCell { pub(crate) fn new( model: String, reasoning_effort: Option, + show_fast_status: bool, directory: PathBuf, version: &'static str, ) -> Self { @@ -1157,6 +1161,7 @@ impl SessionHeaderHistoryCell { model, Style::default(), reasoning_effort, + show_fast_status, directory, version, ) @@ -1166,6 +1171,7 @@ impl SessionHeaderHistoryCell { model: String, model_style: Style, reasoning_effort: Option, + show_fast_status: bool, directory: PathBuf, version: &'static str, ) -> Self { @@ -1174,6 +1180,7 @@ impl SessionHeaderHistoryCell { model, model_style, reasoning_effort, + show_fast_status, directory, } } @@ -1253,6 +1260,10 @@ impl HistoryCell for SessionHeaderHistoryCell { spans.push(Span::from(" ")); spans.push(Span::from(reasoning)); } + if self.show_fast_status { + spans.push(" ".into()); + spans.push(Span::styled("fast", self.model_style.magenta())); + } spans.push(" ".dim()); spans.push(CHANGE_MODEL_HINT_COMMAND.cyan()); spans.push(CHANGE_MODEL_HINT_EXPLANATION.dim()); @@ -2612,6 +2623,7 @@ mod tests { false, Some("Model just became available".to_string()), Some(PlanType::Free), + false, ); let rendered = render_transcript(&cell).join("\n"); @@ -2629,6 +2641,7 @@ mod tests { false, Some("Model just became available".to_string()), Some(PlanType::Free), + false, ); let rendered = render_transcript(&cell).join("\n"); @@ -2645,6 +2658,7 @@ mod tests { true, Some("Model just became available".to_string()), Some(PlanType::Free), + false, ); let rendered = render_transcript(&cell).join("\n"); @@ -2663,6 +2677,7 @@ mod tests { false, Some("Model just became available".to_string()), Some(PlanType::Free), + false, ); let rendered = render_transcript(&cell).join("\n"); @@ -3295,18 +3310,39 @@ mod tests { let cell = SessionHeaderHistoryCell::new( "gpt-4o".to_string(), Some(ReasoningEffortConfig::High), + true, std::env::temp_dir(), "test", ); let lines = render_lines(&cell.display_lines(80)); let model_line = lines - .into_iter() + .iter() + .find(|line| line.contains("model:")) + .expect("model line"); + + assert!(model_line.contains("gpt-4o high fast")); + assert!(model_line.contains("/model to change")); + } + + #[test] + fn session_header_hides_fast_status_when_disabled() { + let cell = SessionHeaderHistoryCell::new( + "gpt-4o".to_string(), + Some(ReasoningEffortConfig::High), + false, + std::env::temp_dir(), + "test", + ); + + let lines = render_lines(&cell.display_lines(80)); + let model_line = lines + .iter() .find(|line| line.contains("model:")) .expect("model line"); assert!(model_line.contains("gpt-4o high")); - assert!(model_line.contains("/model to change")); + assert!(!model_line.contains("fast")); } #[test] diff --git a/codex-rs/tui/src/snapshots/codex_tui__app__tests__clear_ui_after_long_transcript_fresh_header_only.snap b/codex-rs/tui/src/snapshots/codex_tui__app__tests__clear_ui_after_long_transcript_fresh_header_only.snap index abb423ff1f..98d6015064 100644 --- a/codex-rs/tui/src/snapshots/codex_tui__app__tests__clear_ui_after_long_transcript_fresh_header_only.snap +++ b/codex-rs/tui/src/snapshots/codex_tui__app__tests__clear_ui_after_long_transcript_fresh_header_only.snap @@ -1,6 +1,5 @@ --- source: tui/src/app.rs -assertion_line: 3452 expression: rendered --- ╭─────────────────────────────────────────────╮ diff --git a/codex-rs/tui/src/snapshots/codex_tui__history_cell__tests__session_info_availability_nux_tooltip_snapshot.snap b/codex-rs/tui/src/snapshots/codex_tui__history_cell__tests__session_info_availability_nux_tooltip_snapshot.snap index 82869c0171..ad0fc4be1c 100644 --- a/codex-rs/tui/src/snapshots/codex_tui__history_cell__tests__session_info_availability_nux_tooltip_snapshot.snap +++ b/codex-rs/tui/src/snapshots/codex_tui__history_cell__tests__session_info_availability_nux_tooltip_snapshot.snap @@ -1,6 +1,5 @@ --- source: tui/src/history_cell.rs -assertion_line: 2608 expression: rendered --- ╭─────────────────────────────────────╮ From b4cb989563e68639c2bdc748c20c965ba1d830ea Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 5 Mar 2026 00:55:12 -0800 Subject: [PATCH 3/4] refactor: prepare unified exec for zsh-fork backend (#13392) ## Why `shell_zsh_fork` already provides stronger guarantees around which executables receive elevated permissions. To reuse that machinery from unified exec without pushing Unix-specific escalation details through generic runtime code, the escalation bootstrap and session lifetime handling need a cleaner boundary. That boundary also needs to be safe for long-lived sessions: when an intercepted shell session is closed or pruned, any in-flight approval workers and any already-approved escalated child they spawned must be torn down with the session, and the inherited escalation socket must not leak into unrelated subprocesses. ## What Changed - Extracted a reusable `EscalationSession` and `EscalateServer::start_session(...)` in `shell-escalation` so callers can get the wrapper/socket env overlay and keep the escalation server alive without immediately running a one-shot command. - Documented that `EscalationSession::env()` and `ShellCommandExecutor::run(...)` exchange only that env overlay, which callers must merge into their own base shell environment. - Clarified the prepared-exec helper boundary in `core` by naming the new helper APIs around `ExecRequest`, while keeping the legacy `execute_env(...)` entrypoints as thin compatibility wrappers for existing callers that still use the older naming. - Added a small post-spawn hook on the prepared execution path so the parent copy of the inheritable escalation socket is closed immediately after both the existing one-shot shell-command spawn and the unified-exec spawn. - Made session teardown explicit with session-scoped cancellation: dropping an `EscalationSession` or canceling its parent request now stops intercept workers, and the server-spawned escalated child uses `kill_on_drop(true)` so teardown cannot orphan an already-approved child. - Added `UnifiedExecBackendConfig` plumbing through `ToolsConfig`, a `shell::zsh_fork_backend` facade, and an opaque unified-exec spawn-lifecycle hook so unified exec can prepare a wrapped `zsh -c/-lc` request without storing `EscalationSession` directly in generic process/runtime code. - Kept the existing `shell_command` zsh-fork behavior intact on top of the new bootstrap path. Tool selection is unchanged in this PR: when `shell_zsh_fork` is enabled, `ShellCommand` still wins over `exec_command`. ## Verification - `cargo test -p codex-shell-escalation` - includes coverage for `start_session_exposes_wrapper_env_overlay` - includes coverage for `exec_closes_parent_socket_after_shell_spawn` - includes coverage for `dropping_session_aborts_intercept_workers_and_kills_spawned_child` - `cargo test -p codex-core shell_zsh_fork_prefers_shell_command_over_unified_exec` - `cargo test -p codex-core --test all shell_zsh_fork_prompts_for_skill_script_execution` --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/openai/codex/pull/13392). * #13432 * __->__ #13392 --- codex-rs/core/src/exec.rs | 14 +- codex-rs/core/src/sandboxing/mod.rs | 17 +- codex-rs/core/src/tasks/user_shell.rs | 4 +- codex-rs/core/src/tools/runtimes/shell.rs | 7 +- .../tools/runtimes/shell/unix_escalation.rs | 133 ++++- .../runtimes/shell/unix_escalation_tests.rs | 4 + .../tools/runtimes/shell/zsh_fork_backend.rs | 115 +++++ .../core/src/tools/runtimes/unified_exec.rs | 51 +- codex-rs/core/src/tools/spec.rs | 18 + codex-rs/core/src/unified_exec/mod.rs | 4 + codex-rs/core/src/unified_exec/process.rs | 17 +- .../core/src/unified_exec/process_manager.rs | 8 +- codex-rs/shell-escalation/src/lib.rs | 2 + .../src/unix/escalate_server.rs | 470 ++++++++++++++++-- codex-rs/shell-escalation/src/unix/mod.rs | 1 + .../shell-escalation/src/unix/stopwatch.rs | 32 +- 16 files changed, 821 insertions(+), 76 deletions(-) create mode 100644 codex-rs/core/src/tools/runtimes/shell/zsh_fork_backend.rs diff --git a/codex-rs/core/src/exec.rs b/codex-rs/core/src/exec.rs index 01c559c1c5..8779b2e1c3 100644 --- a/codex-rs/core/src/exec.rs +++ b/codex-rs/core/src/exec.rs @@ -231,10 +231,11 @@ pub async fn process_exec_tool_call( crate::sandboxing::execute_env(exec_req, stdout_stream).await } -pub(crate) async fn execute_exec_env( - env: ExecRequest, +pub(crate) async fn execute_exec_request( + exec_request: ExecRequest, sandbox_policy: &SandboxPolicy, stdout_stream: Option, + after_spawn: Option>, ) -> Result { let ExecRequest { command, @@ -248,7 +249,7 @@ pub(crate) async fn execute_exec_env( sandbox_policy: _sandbox_policy_from_env, justification, arg0, - } = env; + } = exec_request; let params = ExecParams { command, @@ -263,7 +264,7 @@ pub(crate) async fn execute_exec_env( }; let start = Instant::now(); - let raw_output_result = exec(params, sandbox, sandbox_policy, stdout_stream).await; + let raw_output_result = exec(params, sandbox, sandbox_policy, stdout_stream, after_spawn).await; let duration = start.elapsed(); finalize_exec_result(raw_output_result, sandbox, duration) } @@ -693,6 +694,7 @@ async fn exec( sandbox: SandboxType, sandbox_policy: &SandboxPolicy, stdout_stream: Option, + after_spawn: Option>, ) -> Result { #[cfg(target_os = "windows")] if sandbox == SandboxType::WindowsRestrictedToken @@ -738,6 +740,9 @@ async fn exec( env, }) .await?; + if let Some(after_spawn) = after_spawn { + after_spawn(); + } consume_truncated_output(child, expiration, stdout_stream).await } @@ -1136,6 +1141,7 @@ mod tests { SandboxType::None, &SandboxPolicy::new_read_only_policy(), None, + None, ) .await?; assert!(output.timed_out); diff --git a/codex-rs/core/src/sandboxing/mod.rs b/codex-rs/core/src/sandboxing/mod.rs index 10c5d71da2..8cbb18e908 100644 --- a/codex-rs/core/src/sandboxing/mod.rs +++ b/codex-rs/core/src/sandboxing/mod.rs @@ -10,7 +10,7 @@ use crate::exec::ExecExpiration; use crate::exec::ExecToolCallOutput; use crate::exec::SandboxType; use crate::exec::StdoutStream; -use crate::exec::execute_exec_env; +use crate::exec::execute_exec_request; use crate::landlock::allow_network_for_proxy; use crate::landlock::create_linux_sandbox_command_args; use crate::protocol::SandboxPolicy; @@ -418,11 +418,20 @@ impl SandboxManager { } pub async fn execute_env( - env: ExecRequest, + exec_request: ExecRequest, stdout_stream: Option, ) -> crate::error::Result { - let effective_policy = env.sandbox_policy.clone(); - execute_exec_env(env, &effective_policy, stdout_stream).await + let effective_policy = exec_request.sandbox_policy.clone(); + execute_exec_request(exec_request, &effective_policy, stdout_stream, None).await +} + +pub async fn execute_exec_request_with_after_spawn( + exec_request: ExecRequest, + stdout_stream: Option, + after_spawn: Option>, +) -> crate::error::Result { + let effective_policy = exec_request.sandbox_policy.clone(); + execute_exec_request(exec_request, &effective_policy, stdout_stream, after_spawn).await } #[cfg(test)] diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index deab56bc5a..2f77d9fcee 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -14,7 +14,7 @@ use crate::exec::ExecToolCallOutput; use crate::exec::SandboxType; use crate::exec::StdoutStream; use crate::exec::StreamOutput; -use crate::exec::execute_exec_env; +use crate::exec::execute_exec_request; use crate::exec_env::create_env; use crate::parse_command::parse_command; use crate::protocol::EventMsg; @@ -177,7 +177,7 @@ pub(crate) async fn execute_user_shell_command( tx_event: session.get_tx_event(), }); - let exec_result = execute_exec_env(exec_env, &sandbox_policy, stdout_stream) + let exec_result = execute_exec_request(exec_env, &sandbox_policy, stdout_stream, None) .or_cancel(&cancellation_token) .await; diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index 101d9309fd..90a4651428 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -5,7 +5,8 @@ Executes shell requests under the orchestrator: asks for approval when needed, builds a CommandSpec, and runs it under the current SandboxAttempt. */ #[cfg(unix)] -mod unix_escalation; +pub(crate) mod unix_escalation; +pub(crate) mod zsh_fork_backend; use crate::command_canonicalization::canonicalize_command_for_approval; use crate::exec::ExecToolCallOutput; @@ -80,7 +81,6 @@ pub(crate) enum ShellRuntimeBackend { #[derive(Default)] pub struct ShellRuntime { - #[cfg_attr(not(unix), allow(dead_code))] backend: ShellRuntimeBackend, } @@ -215,9 +215,8 @@ impl ToolRuntime for ShellRuntime { command }; - #[cfg(unix)] if self.backend == ShellRuntimeBackend::ShellCommandZshFork { - match unix_escalation::try_run_zsh_fork(req, attempt, ctx, &command).await? { + match zsh_fork_backend::maybe_run_shell_command(req, attempt, ctx, &command).await? { Some(out) => return Ok(out), None => { tracing::warn!( diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs index 36f63be4fa..1be2654cec 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -6,6 +6,7 @@ use crate::exec::ExecToolCallOutput; use crate::exec::SandboxType; use crate::exec::is_likely_sandbox_denied; use crate::features::Feature; +use crate::sandboxing::ExecRequest; use crate::sandboxing::SandboxPermissions; use crate::shell::ShellType; use crate::skills::SkillMetadata; @@ -36,6 +37,7 @@ use codex_shell_escalation::EscalationDecision; use codex_shell_escalation::EscalationExecution; use codex_shell_escalation::EscalationPermissions; use codex_shell_escalation::EscalationPolicy; +use codex_shell_escalation::EscalationSession; use codex_shell_escalation::ExecParams; use codex_shell_escalation::ExecResult; use codex_shell_escalation::Permissions as EscalatedPermissions; @@ -51,6 +53,11 @@ use tokio::sync::RwLock; use tokio_util::sync::CancellationToken; use uuid::Uuid; +pub(crate) struct PreparedUnifiedExecZshFork { + pub(crate) exec_request: ExecRequest, + pub(crate) escalation_session: EscalationSession, +} + pub(super) async fn try_run_zsh_fork( req: &ShellRequest, attempt: &SandboxAttempt<'_>, @@ -95,7 +102,7 @@ pub(super) async fn try_run_zsh_fork( justification, arg0, } = sandbox_exec_request; - let ParsedShellCommand { script, login } = extract_shell_script(&command)?; + let ParsedShellCommand { script, login, .. } = extract_shell_script(&command)?; let effective_timeout = Duration::from_millis( req.timeout_ms .unwrap_or(crate::exec::DEFAULT_EXEC_COMMAND_TIMEOUT_MS), @@ -172,6 +179,103 @@ pub(super) async fn try_run_zsh_fork( map_exec_result(attempt.sandbox, exec_result).map(Some) } +pub(crate) async fn prepare_unified_exec_zsh_fork( + req: &crate::tools::runtimes::unified_exec::UnifiedExecRequest, + attempt: &SandboxAttempt<'_>, + ctx: &ToolCtx, + exec_request: ExecRequest, +) -> Result, ToolError> { + let Some(shell_zsh_path) = ctx.session.services.shell_zsh_path.as_ref() else { + tracing::warn!("ZshFork backend specified, but shell_zsh_path is not configured."); + return Ok(None); + }; + if !ctx.session.features().enabled(Feature::ShellZshFork) { + tracing::warn!("ZshFork backend specified, but ShellZshFork feature is not enabled."); + return Ok(None); + } + if !matches!(ctx.session.user_shell().shell_type, ShellType::Zsh) { + tracing::warn!("ZshFork backend specified, but user shell is not Zsh."); + return Ok(None); + } + + let parsed = match extract_shell_script(&exec_request.command) { + Ok(parsed) => parsed, + Err(err) => { + tracing::warn!("ZshFork unified exec fallback: {err:?}"); + return Ok(None); + } + }; + if parsed.program != shell_zsh_path.to_string_lossy() { + tracing::warn!( + "ZshFork backend specified, but unified exec command targets `{}` instead of `{}`.", + parsed.program, + shell_zsh_path.display(), + ); + return Ok(None); + } + + let exec_policy = Arc::new(RwLock::new( + ctx.session.services.exec_policy.current().as_ref().clone(), + )); + let command_executor = CoreShellCommandExecutor { + command: exec_request.command.clone(), + cwd: exec_request.cwd.clone(), + sandbox_policy: exec_request.sandbox_policy.clone(), + sandbox: exec_request.sandbox, + env: exec_request.env.clone(), + network: exec_request.network.clone(), + windows_sandbox_level: exec_request.windows_sandbox_level, + sandbox_permissions: exec_request.sandbox_permissions, + justification: exec_request.justification.clone(), + arg0: exec_request.arg0.clone(), + sandbox_policy_cwd: ctx.turn.cwd.clone(), + macos_seatbelt_profile_extensions: ctx + .turn + .config + .permissions + .macos_seatbelt_profile_extensions + .clone(), + codex_linux_sandbox_exe: ctx.turn.codex_linux_sandbox_exe.clone(), + use_linux_sandbox_bwrap: ctx.turn.features.enabled(Feature::UseLinuxSandboxBwrap), + }; + let main_execve_wrapper_exe = ctx + .session + .services + .main_execve_wrapper_exe + .clone() + .ok_or_else(|| { + ToolError::Rejected( + "zsh fork feature enabled, but execve wrapper is not configured".to_string(), + ) + })?; + let escalation_policy = CoreShellActionProvider { + policy: Arc::clone(&exec_policy), + session: Arc::clone(&ctx.session), + turn: Arc::clone(&ctx.turn), + call_id: ctx.call_id.clone(), + approval_policy: ctx.turn.approval_policy.value(), + sandbox_policy: attempt.policy.clone(), + sandbox_permissions: req.sandbox_permissions, + prompt_permissions: req.additional_permissions.clone(), + stopwatch: Stopwatch::unlimited(), + }; + + let escalate_server = EscalateServer::new( + shell_zsh_path.clone(), + main_execve_wrapper_exe, + escalation_policy, + ); + let escalation_session = escalate_server + .start_session(CancellationToken::new(), Arc::new(command_executor)) + .map_err(|err| ToolError::Rejected(err.to_string()))?; + let mut exec_request = exec_request; + exec_request.env.extend(escalation_session.env().clone()); + Ok(Some(PreparedUnifiedExecZshFork { + exec_request, + escalation_session, + })) +} + struct CoreShellActionProvider { policy: Arc>, session: Arc, @@ -648,17 +752,20 @@ impl ShellCommandExecutor for CoreShellCommandExecutor { &self, _command: Vec, _cwd: PathBuf, - env: HashMap, + env_overlay: HashMap, cancel_rx: CancellationToken, + after_spawn: Option>, ) -> anyhow::Result { let mut exec_env = self.env.clone(); + // `env_overlay` comes from `EscalationSession::env()`, so merge only the + // wrapper/socket variables into the base shell environment. for var in ["CODEX_ESCALATE_SOCKET", "EXEC_WRAPPER", "BASH_EXEC_WRAPPER"] { - if let Some(value) = env.get(var) { + if let Some(value) = env_overlay.get(var) { exec_env.insert(var.to_string(), value.clone()); } } - let result = crate::sandboxing::execute_env( + let result = crate::sandboxing::execute_exec_request_with_after_spawn( crate::sandboxing::ExecRequest { command: self.command.clone(), cwd: self.cwd.clone(), @@ -673,6 +780,7 @@ impl ShellCommandExecutor for CoreShellCommandExecutor { arg0: self.arg0.clone(), }, None, + after_spawn, ) .await?; @@ -809,6 +917,7 @@ impl CoreShellCommandExecutor { #[derive(Debug, Eq, PartialEq)] struct ParsedShellCommand { + program: String, script: String, login: bool, } @@ -817,12 +926,20 @@ fn extract_shell_script(command: &[String]) -> Result Some((script.to_owned(), false)), - [_, flag, script] if flag == "-lc" => Some((script.to_owned(), true)), + if let Some((program, script, login)) = command.windows(3).find_map(|parts| match parts { + [program, flag, script] if flag == "-c" => { + Some((program.to_owned(), script.to_owned(), false)) + } + [program, flag, script] if flag == "-lc" => { + Some((program.to_owned(), script.to_owned(), true)) + } _ => None, }) { - return Ok(ParsedShellCommand { script, login }); + return Ok(ParsedShellCommand { + program, + script, + login, + }); } Err(ToolError::Rejected( diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs index b88f5b0542..ad663a3fed 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs @@ -64,6 +64,7 @@ fn extract_shell_script_preserves_login_flag() { assert_eq!( extract_shell_script(&["/bin/zsh".into(), "-lc".into(), "echo hi".into()]).unwrap(), ParsedShellCommand { + program: "/bin/zsh".to_string(), script: "echo hi".to_string(), login: true, } @@ -71,6 +72,7 @@ fn extract_shell_script_preserves_login_flag() { assert_eq!( extract_shell_script(&["/bin/zsh".into(), "-c".into(), "echo hi".into()]).unwrap(), ParsedShellCommand { + program: "/bin/zsh".to_string(), script: "echo hi".to_string(), login: false, } @@ -89,6 +91,7 @@ fn extract_shell_script_supports_wrapped_command_prefixes() { ]) .unwrap(), ParsedShellCommand { + program: "/bin/zsh".to_string(), script: "echo hello".to_string(), login: true, } @@ -105,6 +108,7 @@ fn extract_shell_script_supports_wrapped_command_prefixes() { ]) .unwrap(), ParsedShellCommand { + program: "/bin/zsh".to_string(), script: "pwd".to_string(), login: false, } diff --git a/codex-rs/core/src/tools/runtimes/shell/zsh_fork_backend.rs b/codex-rs/core/src/tools/runtimes/shell/zsh_fork_backend.rs new file mode 100644 index 0000000000..e88f0caa74 --- /dev/null +++ b/codex-rs/core/src/tools/runtimes/shell/zsh_fork_backend.rs @@ -0,0 +1,115 @@ +use super::ShellRequest; +use crate::exec::ExecToolCallOutput; +use crate::sandboxing::ExecRequest; +use crate::tools::runtimes::unified_exec::UnifiedExecRequest; +use crate::tools::sandboxing::SandboxAttempt; +use crate::tools::sandboxing::ToolCtx; +use crate::tools::sandboxing::ToolError; +use crate::unified_exec::SpawnLifecycleHandle; + +pub(crate) struct PreparedUnifiedExecSpawn { + pub(crate) exec_request: ExecRequest, + pub(crate) spawn_lifecycle: SpawnLifecycleHandle, +} + +/// Runs the zsh-fork shell-command backend when this request should be handled +/// by executable-level escalation instead of the default shell runtime. +/// +/// Returns `Ok(None)` when the current platform or request shape should fall +/// back to the normal shell-command path. +pub(crate) async fn maybe_run_shell_command( + req: &ShellRequest, + attempt: &SandboxAttempt<'_>, + ctx: &ToolCtx, + command: &[String], +) -> Result, ToolError> { + imp::maybe_run_shell_command(req, attempt, ctx, command).await +} + +/// Prepares unified exec to launch through the zsh-fork backend when the +/// request matches a wrapped `zsh -c/-lc` command on a supported platform. +/// +/// Returns the transformed `ExecRequest` plus a spawn lifecycle that keeps the +/// escalation server alive for the session and performs post-spawn cleanup. +/// Returns `Ok(None)` when unified exec should use its normal spawn path. +pub(crate) async fn maybe_prepare_unified_exec( + req: &UnifiedExecRequest, + attempt: &SandboxAttempt<'_>, + ctx: &ToolCtx, + exec_request: ExecRequest, +) -> Result, ToolError> { + imp::maybe_prepare_unified_exec(req, attempt, ctx, exec_request).await +} + +#[cfg(unix)] +mod imp { + use super::*; + use crate::tools::runtimes::shell::unix_escalation; + use crate::unified_exec::SpawnLifecycle; + use codex_shell_escalation::EscalationSession; + + #[derive(Debug)] + struct ZshForkSpawnLifecycle { + escalation_session: EscalationSession, + } + + impl SpawnLifecycle for ZshForkSpawnLifecycle { + fn after_spawn(&mut self) { + self.escalation_session.close_client_socket(); + } + } + + pub(super) async fn maybe_run_shell_command( + req: &ShellRequest, + attempt: &SandboxAttempt<'_>, + ctx: &ToolCtx, + command: &[String], + ) -> Result, ToolError> { + unix_escalation::try_run_zsh_fork(req, attempt, ctx, command).await + } + + pub(super) async fn maybe_prepare_unified_exec( + req: &UnifiedExecRequest, + attempt: &SandboxAttempt<'_>, + ctx: &ToolCtx, + exec_request: ExecRequest, + ) -> Result, ToolError> { + let Some(prepared) = + unix_escalation::prepare_unified_exec_zsh_fork(req, attempt, ctx, exec_request).await? + else { + return Ok(None); + }; + + Ok(Some(PreparedUnifiedExecSpawn { + exec_request: prepared.exec_request, + spawn_lifecycle: Box::new(ZshForkSpawnLifecycle { + escalation_session: prepared.escalation_session, + }), + })) + } +} + +#[cfg(not(unix))] +mod imp { + use super::*; + + pub(super) async fn maybe_run_shell_command( + req: &ShellRequest, + attempt: &SandboxAttempt<'_>, + ctx: &ToolCtx, + command: &[String], + ) -> Result, ToolError> { + let _ = (req, attempt, ctx, command); + Ok(None) + } + + pub(super) async fn maybe_prepare_unified_exec( + req: &UnifiedExecRequest, + attempt: &SandboxAttempt<'_>, + ctx: &ToolCtx, + exec_request: ExecRequest, + ) -> Result, ToolError> { + let _ = (req, attempt, ctx, exec_request); + Ok(None) + } +} diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index f4d823bf0e..d557a4340b 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -16,6 +16,7 @@ use crate::tools::network_approval::NetworkApprovalMode; use crate::tools::network_approval::NetworkApprovalSpec; use crate::tools::runtimes::build_command_spec; use crate::tools::runtimes::maybe_wrap_shell_lc_with_snapshot; +use crate::tools::runtimes::shell::zsh_fork_backend; use crate::tools::sandboxing::Approvable; use crate::tools::sandboxing::ApprovalCtx; use crate::tools::sandboxing::ExecApprovalRequirement; @@ -28,6 +29,8 @@ use crate::tools::sandboxing::ToolError; use crate::tools::sandboxing::ToolRuntime; use crate::tools::sandboxing::sandbox_override_for_first_attempt; use crate::tools::sandboxing::with_cached_approval; +use crate::tools::spec::UnifiedExecBackendConfig; +use crate::unified_exec::NoopSpawnLifecycle; use crate::unified_exec::UnifiedExecError; use crate::unified_exec::UnifiedExecProcess; use crate::unified_exec::UnifiedExecProcessManager; @@ -63,11 +66,12 @@ pub struct UnifiedExecApprovalKey { pub struct UnifiedExecRuntime<'a> { manager: &'a UnifiedExecProcessManager, + backend: UnifiedExecBackendConfig, } impl<'a> UnifiedExecRuntime<'a> { - pub fn new(manager: &'a UnifiedExecProcessManager) -> Self { - Self { manager } + pub fn new(manager: &'a UnifiedExecProcessManager, backend: UnifiedExecBackendConfig) -> Self { + Self { manager, backend } } } @@ -184,6 +188,47 @@ impl<'a> ToolRuntime for UnifiedExecRunt if let Some(network) = req.network.as_ref() { network.apply_to_env(&mut env); } + if self.backend == UnifiedExecBackendConfig::ZshFork { + let spec = build_command_spec( + &command, + &req.cwd, + &env, + ExecExpiration::DefaultTimeout, + req.sandbox_permissions, + req.additional_permissions.clone(), + req.justification.clone(), + ) + .map_err(|_| ToolError::Rejected("missing command line for PTY".to_string()))?; + let exec_env = attempt + .env_for(spec, req.network.as_ref()) + .map_err(|err| ToolError::Codex(err.into()))?; + match zsh_fork_backend::maybe_prepare_unified_exec(req, attempt, ctx, exec_env).await? { + Some(prepared) => { + return self + .manager + .open_session_with_exec_env( + &prepared.exec_request, + req.tty, + prepared.spawn_lifecycle, + ) + .await + .map_err(|err| match err { + UnifiedExecError::SandboxDenied { output, .. } => { + ToolError::Codex(CodexErr::Sandbox(SandboxErr::Denied { + output: Box::new(output), + network_policy_decision: None, + })) + } + other => ToolError::Rejected(other.to_string()), + }); + } + None => { + tracing::warn!( + "UnifiedExec ZshFork backend specified, but conditions for using it were not met, falling back to direct execution", + ); + } + } + } let spec = build_command_spec( &command, &req.cwd, @@ -198,7 +243,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt .env_for(spec, req.network.as_ref()) .map_err(|err| ToolError::Codex(err.into()))?; self.manager - .open_session_with_exec_env(&exec_env, req.tty) + .open_session_with_exec_env(&exec_env, req.tty, Box::new(NoopSpawnLifecycle)) .await .map_err(|err| match err { UnifiedExecError::SandboxDenied { output, .. } => { diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 00867b20d8..cabcf96cf9 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -44,10 +44,17 @@ pub enum ShellCommandBackendConfig { ZshFork, } +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub enum UnifiedExecBackendConfig { + Direct, + ZshFork, +} + #[derive(Debug, Clone)] pub(crate) struct ToolsConfig { pub shell_type: ConfigShellToolType, shell_command_backend: ShellCommandBackendConfig, + pub unified_exec_backend: UnifiedExecBackendConfig, pub allow_login_shell: bool, pub apply_patch_tool_type: Option, pub web_search_mode: Option, @@ -102,6 +109,12 @@ impl ToolsConfig { } else { ShellCommandBackendConfig::Classic }; + let unified_exec_backend = + if features.enabled(Feature::ShellTool) && features.enabled(Feature::ShellZshFork) { + UnifiedExecBackendConfig::ZshFork + } else { + UnifiedExecBackendConfig::Direct + }; let shell_type = if !features.enabled(Feature::ShellTool) { ConfigShellToolType::Disabled @@ -140,6 +153,7 @@ impl ToolsConfig { Self { shell_type, shell_command_backend, + unified_exec_backend, allow_login_shell: true, apply_patch_tool_type, web_search_mode: *web_search_mode, @@ -2817,6 +2831,10 @@ mod tests { tools_config.shell_command_backend, ShellCommandBackendConfig::ZshFork ); + assert_eq!( + tools_config.unified_exec_backend, + UnifiedExecBackendConfig::ZshFork + ); } #[test] diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 1d949dc2b9..6323857150 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -49,6 +49,10 @@ pub(crate) fn set_deterministic_process_ids_for_tests(enabled: bool) { } pub(crate) use errors::UnifiedExecError; +pub(crate) use process::NoopSpawnLifecycle; +#[cfg(unix)] +pub(crate) use process::SpawnLifecycle; +pub(crate) use process::SpawnLifecycleHandle; pub(crate) use process::UnifiedExecProcess; pub(crate) const MIN_YIELD_TIME_MS: u64 = 250; diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index 69790dbad2..ba5352ff69 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -24,6 +24,17 @@ use super::UNIFIED_EXEC_OUTPUT_MAX_TOKENS; use super::UnifiedExecError; use super::head_tail_buffer::HeadTailBuffer; +pub(crate) trait SpawnLifecycle: std::fmt::Debug + Send + Sync { + fn after_spawn(&mut self) {} +} + +pub(crate) type SpawnLifecycleHandle = Box; + +#[derive(Debug, Default)] +pub(crate) struct NoopSpawnLifecycle; + +impl SpawnLifecycle for NoopSpawnLifecycle {} + pub(crate) type OutputBuffer = Arc>; pub(crate) struct OutputHandles { pub(crate) output_buffer: OutputBuffer, @@ -44,6 +55,7 @@ pub(crate) struct UnifiedExecProcess { output_drained: Arc, output_task: JoinHandle<()>, sandbox_type: SandboxType, + _spawn_lifecycle: SpawnLifecycleHandle, } impl UnifiedExecProcess { @@ -51,6 +63,7 @@ impl UnifiedExecProcess { process_handle: ExecCommandSession, initial_output_rx: tokio::sync::broadcast::Receiver>, sandbox_type: SandboxType, + spawn_lifecycle: SpawnLifecycleHandle, ) -> Self { let output_buffer = Arc::new(Mutex::new(HeadTailBuffer::default())); let output_notify = Arc::new(Notify::new()); @@ -92,6 +105,7 @@ impl UnifiedExecProcess { output_drained, output_task, sandbox_type, + _spawn_lifecycle: spawn_lifecycle, } } @@ -196,13 +210,14 @@ impl UnifiedExecProcess { pub(super) async fn from_spawned( spawned: SpawnedPty, sandbox_type: SandboxType, + spawn_lifecycle: SpawnLifecycleHandle, ) -> Result { let SpawnedPty { session: process_handle, output_rx, mut exit_rx, } = spawned; - let managed = Self::new(process_handle, output_rx, sandbox_type); + let managed = Self::new(process_handle, output_rx, sandbox_type, spawn_lifecycle); let exit_ready = matches!(exit_rx.try_recv(), Ok(_) | Err(TryRecvError::Closed)); diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 821abc683e..4ba125009e 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -49,6 +49,7 @@ use crate::unified_exec::generate_chunk_id; use crate::unified_exec::head_tail_buffer::HeadTailBuffer; use crate::unified_exec::process::OutputBuffer; use crate::unified_exec::process::OutputHandles; +use crate::unified_exec::process::SpawnLifecycleHandle; use crate::unified_exec::process::UnifiedExecProcess; use crate::unified_exec::resolve_max_tokens; @@ -528,6 +529,7 @@ impl UnifiedExecProcessManager { &self, env: &ExecRequest, tty: bool, + mut spawn_lifecycle: SpawnLifecycleHandle, ) -> Result { let (program, args) = env .command @@ -555,7 +557,8 @@ impl UnifiedExecProcessManager { }; let spawned = spawn_result.map_err(|err| UnifiedExecError::create_process(err.to_string()))?; - UnifiedExecProcess::from_spawned(spawned, env.sandbox).await + spawn_lifecycle.after_spawn(); + UnifiedExecProcess::from_spawned(spawned, env.sandbox, spawn_lifecycle).await } pub(super) async fn open_session_with_sandbox( @@ -569,7 +572,8 @@ impl UnifiedExecProcessManager { Some(context.session.conversation_id), )); let mut orchestrator = ToolOrchestrator::new(); - let mut runtime = UnifiedExecRuntime::new(self); + let mut runtime = + UnifiedExecRuntime::new(self, context.turn.tools_config.unified_exec_backend); let exec_approval_requirement = context .session .services diff --git a/codex-rs/shell-escalation/src/lib.rs b/codex-rs/shell-escalation/src/lib.rs index 1cc42a46db..6d08a78716 100644 --- a/codex-rs/shell-escalation/src/lib.rs +++ b/codex-rs/shell-escalation/src/lib.rs @@ -14,6 +14,8 @@ pub use unix::EscalationPermissions; #[cfg(unix)] pub use unix::EscalationPolicy; #[cfg(unix)] +pub use unix::EscalationSession; +#[cfg(unix)] pub use unix::ExecParams; #[cfg(unix)] pub use unix::ExecResult; diff --git a/codex-rs/shell-escalation/src/unix/escalate_server.rs b/codex-rs/shell-escalation/src/unix/escalate_server.rs index 92d9adb9d3..bdd248a113 100644 --- a/codex-rs/shell-escalation/src/unix/escalate_server.rs +++ b/codex-rs/shell-escalation/src/unix/escalate_server.rs @@ -3,11 +3,14 @@ use std::os::fd::AsRawFd; use std::path::PathBuf; use std::process::Stdio; use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; use anyhow::Context as _; use codex_utils_absolute_path::AbsolutePathBuf; +use socket2::Socket; use tokio::process::Command; +use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use crate::unix::escalate_protocol::ESCALATE_SOCKET_ENV_VAR; @@ -32,12 +35,20 @@ use crate::unix::socket::AsyncSocket; #[async_trait::async_trait] pub trait ShellCommandExecutor: Send + Sync { /// Runs the requested shell command and returns the captured result. + /// + /// `env_overlay` contains only the wrapper/socket variables exported by + /// `EscalationSession::env()`, not a complete child environment. + /// Implementations should merge it into whatever base environment they use + /// for the shell process. `after_spawn` should be invoked immediately after + /// the shell process has been spawned so the parent copy of the inherited + /// escalation socket can be closed. async fn run( &self, command: Vec, cwd: PathBuf, - env: HashMap, + env_overlay: HashMap, cancel_rx: CancellationToken, + after_spawn: Option>, ) -> anyhow::Result; /// Prepares an escalated subcommand for execution on the server side. @@ -82,6 +93,38 @@ pub struct PreparedExec { pub arg0: Option, } +#[derive(Debug)] +pub struct EscalationSession { + env: HashMap, + task: JoinHandle>, + client_socket: Arc>>, + cancellation_token: CancellationToken, +} + +impl EscalationSession { + /// Returns just the environment overlay needed by the execve wrapper. + /// + /// Callers should merge this into their own child-process environment + /// rather than treating it as the full environment for the shell. + pub fn env(&self) -> &HashMap { + &self.env + } + + pub fn close_client_socket(&self) { + if let Ok(mut client_socket) = self.client_socket.lock() { + client_socket.take(); + } + } +} + +impl Drop for EscalationSession { + fn drop(&mut self) { + self.close_client_socket(); + self.cancellation_token.cancel(); + self.task.abort(); + } +} + pub struct EscalateServer { bash_path: PathBuf, execve_wrapper: PathBuf, @@ -106,29 +149,9 @@ impl EscalateServer { cancel_rx: CancellationToken, command_executor: Arc, ) -> anyhow::Result { - let (escalate_server, escalate_client) = AsyncDatagramSocket::pair()?; - let client_socket = escalate_client.into_inner(); - // Only the client endpoint should cross exec into the wrapper process. - client_socket.set_cloexec(false)?; - let escalate_task = tokio::spawn(escalate_task( - escalate_server, - Arc::clone(&self.policy), - Arc::clone(&command_executor), - )); - let mut env = std::env::vars().collect::>(); - env.insert( - ESCALATE_SOCKET_ENV_VAR.to_string(), - client_socket.as_raw_fd().to_string(), - ); - env.insert( - EXEC_WRAPPER_ENV_VAR.to_string(), - self.execve_wrapper.to_string_lossy().to_string(), - ); - env.insert( - LEGACY_BASH_EXEC_WRAPPER_ENV_VAR.to_string(), - self.execve_wrapper.to_string_lossy().to_string(), - ); - + let session = self.start_session(cancel_rx.clone(), Arc::clone(&command_executor))?; + let env_overlay = session.env().clone(); + let client_socket = Arc::clone(&session.client_socket); let command = vec![ self.bash_path.to_string_lossy().to_string(), if params.login == Some(false) { @@ -140,20 +163,80 @@ impl EscalateServer { ]; let workdir = AbsolutePathBuf::try_from(params.workdir)?; let result = command_executor - .run(command, workdir.to_path_buf(), env, cancel_rx) + .run( + command, + workdir.to_path_buf(), + env_overlay, + cancel_rx, + Some(Box::new(move || { + if let Ok(mut client_socket) = client_socket.lock() { + client_socket.take(); + } + })), + ) .await?; - escalate_task.abort(); Ok(result) } + + /// Starts an escalation session and returns the environment overlay a shell + /// needs in order to route intercepted execs through this server. + /// + /// This does not spawn the shell itself. Callers own process creation and + /// only use the returned environment plus the session lifetime handle. + pub fn start_session( + &self, + parent_cancellation_token: CancellationToken, + command_executor: Arc, + ) -> anyhow::Result { + let cancellation_token = CancellationToken::new(); + let (escalate_server, escalate_client) = AsyncDatagramSocket::pair()?; + let client_socket = escalate_client.into_inner(); + let client_socket_fd = client_socket.as_raw_fd(); + // Only the client endpoint should cross exec into the wrapper process. + client_socket.set_cloexec(false)?; + let client_socket = Arc::new(Mutex::new(Some(client_socket))); + let task = tokio::spawn(escalate_task( + escalate_server, + Arc::clone(&self.policy), + Arc::clone(&command_executor), + parent_cancellation_token, + cancellation_token.clone(), + )); + let mut env = HashMap::new(); + env.insert( + ESCALATE_SOCKET_ENV_VAR.to_string(), + client_socket_fd.to_string(), + ); + env.insert( + EXEC_WRAPPER_ENV_VAR.to_string(), + self.execve_wrapper.to_string_lossy().to_string(), + ); + env.insert( + LEGACY_BASH_EXEC_WRAPPER_ENV_VAR.to_string(), + self.execve_wrapper.to_string_lossy().to_string(), + ); + Ok(EscalationSession { + env, + task, + client_socket, + cancellation_token, + }) + } } async fn escalate_task( socket: AsyncDatagramSocket, policy: Arc, command_executor: Arc, + parent_cancellation_token: CancellationToken, + session_cancellation_token: CancellationToken, ) -> anyhow::Result<()> { loop { - let (_, mut fds) = socket.receive_with_fds().await?; + let (_, mut fds) = tokio::select! { + received = socket.receive_with_fds() => received?, + _ = parent_cancellation_token.cancelled() => return Ok(()), + _ = session_cancellation_token.cancelled() => return Ok(()), + }; if fds.len() != 1 { tracing::error!("expected 1 fd in datagram handshake, got {}", fds.len()); continue; @@ -161,9 +244,17 @@ async fn escalate_task( let stream_socket = AsyncSocket::from_fd(fds.remove(0))?; let policy = Arc::clone(&policy); let command_executor = Arc::clone(&command_executor); + let parent_cancellation_token = parent_cancellation_token.clone(); + let session_cancellation_token = session_cancellation_token.clone(); tokio::spawn(async move { - if let Err(err) = - handle_escalate_session_with_policy(stream_socket, policy, command_executor).await + if let Err(err) = handle_escalate_session_with_policy( + stream_socket, + policy, + command_executor, + parent_cancellation_token, + session_cancellation_token, + ) + .await { tracing::error!("escalate session failed: {err:?}"); } @@ -175,18 +266,27 @@ async fn handle_escalate_session_with_policy( socket: AsyncSocket, policy: Arc, command_executor: Arc, + parent_cancellation_token: CancellationToken, + session_cancellation_token: CancellationToken, ) -> anyhow::Result<()> { let EscalateRequest { file, argv, workdir, env, - } = socket.receive::().await?; + } = tokio::select! { + request = socket.receive::() => request?, + _ = parent_cancellation_token.cancelled() => return Ok(()), + _ = session_cancellation_token.cancelled() => return Ok(()), + }; let program = AbsolutePathBuf::resolve_path_against_base(file, workdir.as_path())?; - let decision = policy - .determine_action(&program, &argv, &workdir) - .await - .context("failed to determine escalation action")?; + let decision = tokio::select! { + decision = policy.determine_action(&program, &argv, &workdir) => { + decision.context("failed to determine escalation action")? + } + _ = parent_cancellation_token.cancelled() => return Ok(()), + _ = session_cancellation_token.cancelled() => return Ok(()), + }; tracing::debug!("decided {decision:?} for {program:?} {argv:?} {workdir:?}"); @@ -204,10 +304,13 @@ async fn handle_escalate_session_with_policy( action: EscalateAction::Escalate, }) .await?; - let (msg, fds) = socket - .receive_with_fds::() - .await - .context("failed to receive SuperExecMessage")?; + let (msg, fds) = tokio::select! { + message = socket.receive_with_fds::() => { + message.context("failed to receive SuperExecMessage")? + } + _ = parent_cancellation_token.cancelled() => return Ok(()), + _ = session_cancellation_token.cancelled() => return Ok(()), + }; if fds.len() != msg.fds.len() { return Err(anyhow::anyhow!( "mismatched number of fds in SuperExecMessage: {} in the message, {} from the control message", @@ -231,9 +334,11 @@ async fn handle_escalate_session_with_policy( cwd, env, arg0, - } = command_executor - .prepare_escalated_exec(&program, &argv, &workdir, env, execution) - .await?; + } = tokio::select! { + prepared = command_executor.prepare_escalated_exec(&program, &argv, &workdir, env, execution) => prepared?, + _ = parent_cancellation_token.cancelled() => return Ok(()), + _ = session_cancellation_token.cancelled() => return Ok(()), + }; let (program, args) = command .split_first() .ok_or_else(|| anyhow::anyhow!("prepared escalated command must not be empty"))?; @@ -245,7 +350,8 @@ async fn handle_escalate_session_with_policy( .current_dir(&cwd) .stdin(Stdio::null()) .stdout(Stdio::null()) - .stderr(Stdio::null()); + .stderr(Stdio::null()) + .kill_on_drop(true); unsafe { command.pre_exec(move || { for (dst_fd, src_fd) in msg.fds.iter().zip(&fds) { @@ -255,7 +361,17 @@ async fn handle_escalate_session_with_policy( }); } let mut child = command.spawn()?; - let exit_status = child.wait().await?; + let exit_status = tokio::select! { + status = child.wait() => status?, + _ = parent_cancellation_token.cancelled() => { + let _ = child.start_kill(); + child.wait().await? + } + _ = session_cancellation_token.cancelled() => { + let _ = child.start_kill(); + child.wait().await? + } + }; socket .send(SuperExecResult { exit_code: exit_status.code().unwrap_or(127), @@ -282,7 +398,15 @@ mod tests { use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; use std::collections::HashMap; + use std::os::fd::FromRawFd; use std::path::PathBuf; + use std::sync::LazyLock; + use tempfile::TempDir; + use tokio::time::Instant; + use tokio::time::sleep; + + static ESCALATE_SERVER_TEST_LOCK: LazyLock> = + LazyLock::new(|| tokio::sync::Mutex::new(())); struct DeterministicEscalationPolicy { decision: EscalationDecision, @@ -327,8 +451,9 @@ mod tests { &self, _command: Vec, _cwd: PathBuf, - _env: HashMap, + _env_overlay: HashMap, _cancel_rx: CancellationToken, + _after_spawn: Option>, ) -> anyhow::Result { unreachable!("run() is not used by handle_escalate_session_with_policy() tests") } @@ -362,8 +487,9 @@ mod tests { &self, _command: Vec, _cwd: PathBuf, - _env: HashMap, + _env_overlay: HashMap, _cancel_rx: CancellationToken, + _after_spawn: Option>, ) -> anyhow::Result { unreachable!("run() is not used by handle_escalate_session_with_policy() tests") } @@ -391,8 +517,160 @@ mod tests { } } + async fn wait_for_pid_file(pid_file: &std::path::Path) -> anyhow::Result { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if let Ok(contents) = std::fs::read_to_string(pid_file) { + return Ok(contents.trim().parse()?); + } + if Instant::now() >= deadline { + return Err(anyhow::anyhow!( + "timed out waiting for pid file {}", + pid_file.display() + )); + } + sleep(Duration::from_millis(20)).await; + } + } + + fn process_exists(pid: i32) -> bool { + let rc = unsafe { libc::kill(pid, 0) }; + if rc == 0 { + return true; + } + std::io::Error::last_os_error().raw_os_error() != Some(libc::ESRCH) + } + + struct AfterSpawnAssertingShellCommandExecutor; + + #[async_trait::async_trait] + impl ShellCommandExecutor for AfterSpawnAssertingShellCommandExecutor { + async fn run( + &self, + _command: Vec, + _cwd: PathBuf, + env_overlay: HashMap, + _cancel_rx: CancellationToken, + after_spawn: Option>, + ) -> anyhow::Result { + let socket_fd = env_overlay + .get(ESCALATE_SOCKET_ENV_VAR) + .expect("session should export shell escalation socket") + .parse::()?; + assert_ne!(unsafe { libc::fcntl(socket_fd, libc::F_GETFD) }, -1); + after_spawn.expect("one-shot exec should install an after-spawn hook")(); + assert_eq!(unsafe { libc::fcntl(socket_fd, libc::F_GETFD) }, -1); + Ok(ExecResult { + exit_code: 0, + stdout: String::new(), + stderr: String::new(), + output: String::new(), + duration: Duration::ZERO, + timed_out: false, + }) + } + + async fn prepare_escalated_exec( + &self, + _program: &AbsolutePathBuf, + _argv: &[String], + _workdir: &AbsolutePathBuf, + _env: HashMap, + _execution: EscalationExecution, + ) -> anyhow::Result { + unreachable!("prepare_escalated_exec() is not used by exec() tests") + } + } + + async fn wait_for_process_exit(pid: i32) -> anyhow::Result<()> { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if !process_exists(pid) { + return Ok(()); + } + if Instant::now() >= deadline { + return Err(anyhow::anyhow!("timed out waiting for pid {pid} to exit")); + } + sleep(Duration::from_millis(20)).await; + } + } + + /// Verifies that `start_session()` returns only the wrapper/socket env + /// overlay and does not need to touch the configured shell or wrapper + /// executable paths. + /// + /// The `/bin/bash` and `/tmp/codex-execve-wrapper` values here are + /// intentionally fake sentinels: this test asserts that the paths are + /// copied into the exported environment and that the socket fd stays valid + /// until `close_client_socket()` is called. + #[tokio::test] + async fn start_session_exposes_wrapper_env_overlay() -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; + let execve_wrapper = PathBuf::from("/tmp/codex-execve-wrapper"); + let execve_wrapper_str = execve_wrapper.to_string_lossy().to_string(); + let server = EscalateServer::new( + PathBuf::from("/bin/bash"), + execve_wrapper.clone(), + DeterministicEscalationPolicy { + decision: EscalationDecision::run(), + }, + ); + + let session = server.start_session( + CancellationToken::new(), + Arc::new(ForwardingShellCommandExecutor), + )?; + let env = session.env(); + assert_eq!(env.get(EXEC_WRAPPER_ENV_VAR), Some(&execve_wrapper_str)); + assert_eq!( + env.get(LEGACY_BASH_EXEC_WRAPPER_ENV_VAR), + Some(&execve_wrapper_str) + ); + let socket_fd = env + .get(ESCALATE_SOCKET_ENV_VAR) + .expect("session should export shell escalation socket"); + let socket_fd = socket_fd.parse::()?; + assert!(socket_fd >= 0); + assert_ne!(unsafe { libc::fcntl(socket_fd, libc::F_GETFD) }, -1); + session.close_client_socket(); + assert_eq!(unsafe { libc::fcntl(socket_fd, libc::F_GETFD) }, -1); + + Ok(()) + } + + #[tokio::test] + async fn exec_closes_parent_socket_after_shell_spawn() -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; + let server = EscalateServer::new( + PathBuf::from("/bin/bash"), + PathBuf::from("/tmp/codex-execve-wrapper"), + DeterministicEscalationPolicy { + decision: EscalationDecision::run(), + }, + ); + + let result = server + .exec( + ExecParams { + command: "true".to_string(), + workdir: AbsolutePathBuf::current_dir()? + .to_string_lossy() + .to_string(), + timeout_ms: None, + login: Some(false), + }, + CancellationToken::new(), + Arc::new(AfterSpawnAssertingShellCommandExecutor), + ) + .await?; + assert_eq!(0, result.exit_code); + + Ok(()) + } + #[tokio::test] async fn handle_escalate_session_respects_run_in_sandbox_decision() -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; let (server, client) = AsyncSocket::pair()?; let server_task = tokio::spawn(handle_escalate_session_with_policy( server, @@ -400,6 +678,8 @@ mod tests { decision: EscalationDecision::run(), }), Arc::new(ForwardingShellCommandExecutor), + CancellationToken::new(), + CancellationToken::new(), )); let mut env = HashMap::new(); @@ -430,6 +710,7 @@ mod tests { #[tokio::test] async fn handle_escalate_session_resolves_relative_file_against_request_workdir() -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; let (server, client) = AsyncSocket::pair()?; let tmp = tempfile::TempDir::new()?; let workdir = tmp.path().join("workspace"); @@ -443,6 +724,8 @@ mod tests { expected_workdir: workdir.clone(), }), Arc::new(ForwardingShellCommandExecutor), + CancellationToken::new(), + CancellationToken::new(), )); client @@ -466,6 +749,7 @@ mod tests { #[tokio::test] async fn handle_escalate_session_executes_escalated_command() -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; let (server, client) = AsyncSocket::pair()?; let server_task = tokio::spawn(handle_escalate_session_with_policy( server, @@ -473,6 +757,8 @@ mod tests { decision: EscalationDecision::escalate(EscalationExecution::Unsandboxed), }), Arc::new(ForwardingShellCommandExecutor), + CancellationToken::new(), + CancellationToken::new(), )); client @@ -508,6 +794,7 @@ mod tests { #[tokio::test] async fn handle_escalate_session_passes_permissions_to_executor() -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; let (server, client) = AsyncSocket::pair()?; let server_task = tokio::spawn(handle_escalate_session_with_policy( server, @@ -529,6 +816,8 @@ mod tests { ..Default::default() }), }), + CancellationToken::new(), + CancellationToken::new(), )); client @@ -557,4 +846,95 @@ mod tests { server_task.await? } + + #[tokio::test] + async fn dropping_session_aborts_intercept_workers_and_kills_spawned_child() + -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; + let tmp = TempDir::new()?; + let pid_file = tmp.path().join("escalated-child.pid"); + let pid_file_display = pid_file.display().to_string(); + assert!( + !pid_file_display.contains('\''), + "test temp path should not contain single quotes: {pid_file_display}" + ); + let server = EscalateServer::new( + PathBuf::from("/bin/bash"), + PathBuf::from("/tmp/codex-execve-wrapper"), + DeterministicEscalationPolicy { + decision: EscalationDecision::escalate(EscalationExecution::Unsandboxed), + }, + ); + + let session = server.start_session( + CancellationToken::new(), + Arc::new(ForwardingShellCommandExecutor), + )?; + let socket_fd = session + .env() + .get(ESCALATE_SOCKET_ENV_VAR) + .expect("session should export shell escalation socket") + .parse::()?; + let dup_socket_fd = unsafe { libc::dup(socket_fd) }; + assert!(dup_socket_fd >= 0, "expected dup() to succeed"); + let handshake_client = unsafe { AsyncDatagramSocket::from_raw_fd(dup_socket_fd) }?; + let (server_stream, client_stream) = AsyncSocket::pair()?; + // Keep one local reference to the server end alive until the worker has + // responded once. Without that guard, macOS can observe EOF on the + // client side before the transferred fd is fully servicing the stream. + let server_stream_guard = server_stream.into_inner(); + let dup_server_stream_fd = unsafe { libc::dup(server_stream_guard.as_raw_fd()) }; + assert!( + dup_server_stream_fd >= 0, + "expected dup() of server stream to succeed" + ); + let server_stream_fd = unsafe { std::os::fd::OwnedFd::from_raw_fd(dup_server_stream_fd) }; + handshake_client + .send_with_fds(&[0], &[server_stream_fd]) + .await + .context("failed to send handshake datagram")?; + + client_stream + .send(EscalateRequest { + file: PathBuf::from("/bin/sh"), + argv: vec![ + "sh".to_string(), + "-c".to_string(), + format!("echo $$ > '{pid_file_display}' && exec /bin/sleep 100"), + ], + workdir: AbsolutePathBuf::current_dir()?, + env: HashMap::new(), + }) + .await + .context("failed to send EscalateRequest")?; + + let response = client_stream + .receive::() + .await + .context("failed to receive EscalateResponse")?; + assert_eq!( + EscalateResponse { + action: EscalateAction::Escalate, + }, + response + ); + drop(server_stream_guard); + + client_stream + .send_with_fds(SuperExecMessage { fds: Vec::new() }, &[]) + .await + .context("failed to send SuperExecMessage")?; + + let pid = wait_for_pid_file(&pid_file).await?; + assert!( + process_exists(pid), + "expected spawned child pid {pid} to exist" + ); + + drop(session); + + wait_for_process_exit(pid).await?; + + Ok(()) + } } diff --git a/codex-rs/shell-escalation/src/unix/mod.rs b/codex-rs/shell-escalation/src/unix/mod.rs index 6de12297a4..13b7146ab4 100644 --- a/codex-rs/shell-escalation/src/unix/mod.rs +++ b/codex-rs/shell-escalation/src/unix/mod.rs @@ -66,6 +66,7 @@ pub use self::escalate_protocol::EscalateAction; pub use self::escalate_protocol::EscalationDecision; pub use self::escalate_protocol::EscalationExecution; pub use self::escalate_server::EscalateServer; +pub use self::escalate_server::EscalationSession; pub use self::escalate_server::ExecParams; pub use self::escalate_server::ExecResult; pub use self::escalate_server::PreparedExec; diff --git a/codex-rs/shell-escalation/src/unix/stopwatch.rs b/codex-rs/shell-escalation/src/unix/stopwatch.rs index f564434813..19821d8be2 100644 --- a/codex-rs/shell-escalation/src/unix/stopwatch.rs +++ b/codex-rs/shell-escalation/src/unix/stopwatch.rs @@ -9,7 +9,7 @@ use tokio_util::sync::CancellationToken; #[derive(Clone, Debug)] pub struct Stopwatch { - limit: Duration, + limit: Option, inner: Arc>, notify: Arc, } @@ -30,13 +30,27 @@ impl Stopwatch { active_pauses: 0, })), notify: Arc::new(Notify::new()), - limit, + limit: Some(limit), + } + } + + pub fn unlimited() -> Self { + Self { + inner: Arc::new(Mutex::new(StopwatchState { + elapsed: Duration::ZERO, + running_since: Some(Instant::now()), + active_pauses: 0, + })), + notify: Arc::new(Notify::new()), + limit: None, } } pub fn cancellation_token(&self) -> CancellationToken { - let limit = self.limit; let token = CancellationToken::new(); + let Some(limit) = self.limit else { + return token; + }; let cancel = token.clone(); let inner = Arc::clone(&self.inner); let notify = Arc::clone(&self.notify); @@ -208,4 +222,16 @@ mod tests { // Now the stopwatch should resume and hit the limit shortly after. token.cancelled().await; } + + #[tokio::test] + async fn unlimited_stopwatch_never_cancels() { + let stopwatch = Stopwatch::unlimited(); + let token = stopwatch.cancellation_token(); + + assert!( + timeout(Duration::from_millis(30), token.cancelled()) + .await + .is_err() + ); + } } From 73a84c74478d69cf0b17d07cce9f7fe0260a8c0d Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 5 Mar 2026 00:55:24 -0800 Subject: [PATCH 4/4] refactor: route zsh-fork through unified exec --- .../tests/suite/v2/turn_start_zsh_fork.rs | 236 +++++++++++++++--- .../tools/runtimes/shell/unix_escalation.rs | 65 ++--- codex-rs/core/src/tools/spec.rs | 8 +- codex-rs/core/tests/common/zsh_fork.rs | 23 ++ codex-rs/core/tests/suite/approvals.rs | 153 ++++++++++++ codex-rs/core/tests/suite/skill_approval.rs | 202 +++++++++++++++ 6 files changed, 621 insertions(+), 66 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs b/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs index bdbb283432..203194d616 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs @@ -11,7 +11,6 @@ use app_test_support::McpProcess; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence; use app_test_support::create_mock_responses_server_sequence_unchecked; -use app_test_support::create_shell_command_sse_response; use app_test_support::to_response; use codex_app_server_protocol::CommandAction; use codex_app_server_protocol::CommandExecutionApprovalDecision; @@ -35,9 +34,11 @@ use codex_core::features::Feature; use core_test_support::responses; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; +use serde_json::json; use std::collections::BTreeMap; use std::path::Path; use tempfile::TempDir; +use tokio::time::sleep; use tokio::time::timeout; #[cfg(windows)] @@ -61,10 +62,8 @@ async fn turn_start_shell_zsh_fork_executes_command_v2() -> Result<()> { }; eprintln!("using zsh path for zsh-fork test: {}", zsh_path.display()); - let responses = vec![create_shell_command_sse_response( - vec!["echo".to_string(), "hi".to_string()], - None, - Some(5000), + let responses = vec![create_zsh_fork_exec_command_sse_response( + "echo hi", "call-zsh-fork", )?]; let server = create_mock_responses_server_sequence(responses).await; @@ -74,7 +73,7 @@ async fn turn_start_shell_zsh_fork_executes_command_v2() -> Result<()> { "never", &BTreeMap::from([ (Feature::ShellZshFork, true), - (Feature::UnifiedExec, false), + (Feature::UnifiedExec, true), (Feature::ShellSnapshot, false), ]), &zsh_path, @@ -172,14 +171,8 @@ async fn turn_start_shell_zsh_fork_exec_approval_decline_v2() -> Result<()> { eprintln!("using zsh path for zsh-fork test: {}", zsh_path.display()); let responses = vec![ - create_shell_command_sse_response( - vec![ - "python3".to_string(), - "-c".to_string(), - "print(42)".to_string(), - ], - None, - Some(5000), + create_zsh_fork_exec_command_sse_response( + "python3 -c 'print(42)'", "call-zsh-fork-decline", )?, create_final_assistant_message_sse_response("done")?, @@ -191,7 +184,7 @@ async fn turn_start_shell_zsh_fork_exec_approval_decline_v2() -> Result<()> { "untrusted", &BTreeMap::from([ (Feature::ShellZshFork, true), - (Feature::UnifiedExec, false), + (Feature::UnifiedExec, true), (Feature::ShellSnapshot, false), ]), &zsh_path, @@ -307,14 +300,8 @@ async fn turn_start_shell_zsh_fork_exec_approval_cancel_v2() -> Result<()> { }; eprintln!("using zsh path for zsh-fork test: {}", zsh_path.display()); - let responses = vec![create_shell_command_sse_response( - vec![ - "python3".to_string(), - "-c".to_string(), - "print(42)".to_string(), - ], - None, - Some(5000), + let responses = vec![create_zsh_fork_exec_command_sse_response( + "python3 -c 'print(42)'", "call-zsh-fork-cancel", )?]; let server = create_mock_responses_server_sequence(responses).await; @@ -324,7 +311,7 @@ async fn turn_start_shell_zsh_fork_exec_approval_cancel_v2() -> Result<()> { "untrusted", &BTreeMap::from([ (Feature::ShellZshFork, true), - (Feature::UnifiedExec, false), + (Feature::UnifiedExec, true), (Feature::ShellSnapshot, false), ]), &zsh_path, @@ -422,6 +409,181 @@ async fn turn_start_shell_zsh_fork_exec_approval_cancel_v2() -> Result<()> { Ok(()) } +#[tokio::test] +async fn turn_start_shell_zsh_fork_interrupt_kills_approved_subcommand_v2() -> Result<()> { + skip_if_no_network!(Ok(())); + + let tmp = TempDir::new()?; + let codex_home = tmp.path().join("codex_home"); + std::fs::create_dir(&codex_home)?; + let workspace = tmp.path().join("workspace"); + std::fs::create_dir(&workspace)?; + let pid_file = workspace.join("approved-subcommand.pid"); + let pid_file_display = pid_file.display().to_string(); + assert!( + !pid_file_display.contains('\''), + "test workspace path should not contain single quotes: {pid_file_display}" + ); + + let Some(zsh_path) = find_test_zsh_path()? else { + eprintln!("skipping zsh fork interrupt cleanup test: no zsh executable found"); + return Ok(()); + }; + if !supports_exec_wrapper_intercept(&zsh_path) { + eprintln!( + "skipping zsh fork interrupt cleanup test: zsh does not support EXEC_WRAPPER intercepts ({})", + zsh_path.display() + ); + return Ok(()); + } + let zsh_path_display = zsh_path.display().to_string(); + eprintln!("using zsh path for zsh-fork test: {zsh_path_display}"); + + let shell_command = + format!("/bin/sh -c 'echo $$ > \"{pid_file_display}\" && exec /bin/sleep 100'"); + let tool_call_arguments = serde_json::to_string(&json!({ + "cmd": shell_command, + "yield_time_ms": 30_000, + }))?; + let response = responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call( + "call-zsh-fork-interrupt-cleanup", + "exec_command", + &tool_call_arguments, + ), + responses::ev_completed("resp-1"), + ]); + let no_op_response = responses::sse(vec![ + responses::ev_response_created("resp-2"), + responses::ev_completed("resp-2"), + ]); + let server = + create_mock_responses_server_sequence_unchecked(vec![response, no_op_response]).await; + create_config_toml( + &codex_home, + &server.uri(), + "untrusted", + &BTreeMap::from([ + (Feature::ShellZshFork, true), + (Feature::UnifiedExec, true), + (Feature::ShellSnapshot, false), + ]), + &zsh_path, + )?; + + let mut mcp = create_zsh_test_mcp_process(&codex_home, &workspace).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let start_id = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + cwd: Some(workspace.to_string_lossy().into_owned()), + ..Default::default() + }) + .await?; + let start_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(start_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + + let turn_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![V2UserInput::Text { + text: "run the long-lived command".to_string(), + text_elements: Vec::new(), + }], + cwd: Some(workspace.clone()), + approval_policy: Some(codex_app_server_protocol::AskForApproval::UnlessTrusted), + sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::WorkspaceWrite { + writable_roots: vec![workspace.clone().try_into()?], + read_only_access: codex_app_server_protocol::ReadOnlyAccess::FullAccess, + network_access: false, + exclude_tmpdir_env_var: false, + exclude_slash_tmp: false, + }), + model: Some("mock-model".to_string()), + effort: Some(codex_protocol::openai_models::ReasoningEffort::Medium), + summary: Some(codex_protocol::config_types::ReasoningSummary::Auto), + ..Default::default() + }) + .await?; + let turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_id)), + ) + .await??; + let TurnStartResponse { turn } = to_response::(turn_resp)?; + + let mut saw_target_approval = false; + while !saw_target_approval { + let server_req = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let ServerRequest::CommandExecutionRequestApproval { request_id, params } = server_req + else { + panic!("expected CommandExecutionRequestApproval request"); + }; + let approval_command = params.command.clone().unwrap_or_default(); + saw_target_approval = approval_command.contains("/bin/sh") + && approval_command.contains(&pid_file_display) + && !approval_command.contains(&zsh_path_display); + mcp.send_response( + request_id, + serde_json::to_value(CommandExecutionRequestApprovalResponse { + decision: CommandExecutionApprovalDecision::Accept, + })?, + ) + .await?; + } + + let pid = timeout(DEFAULT_READ_TIMEOUT, async { + loop { + if let Ok(contents) = std::fs::read_to_string(&pid_file) { + return Ok::(contents.trim().parse()?); + } + sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await??; + let still_running = std::process::Command::new("/bin/kill") + .args(["-0", &pid.to_string()]) + .status()? + .success(); + assert!( + still_running, + "expected approved intercepted subprocess pid {pid} to be running before interrupt" + ); + + mcp.interrupt_turn_and_wait_for_aborted( + thread.id.clone(), + turn.id.clone(), + DEFAULT_READ_TIMEOUT, + ) + .await?; + + timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let still_running = std::process::Command::new("/bin/kill") + .args(["-0", &pid.to_string()]) + .status()? + .success(); + if !still_running { + return Ok::<(), anyhow::Error>(()); + } + sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await??; + + Ok(()) +} + #[tokio::test] async fn turn_start_shell_zsh_fork_subcommand_decline_marks_parent_declined_v2() -> Result<()> { skip_if_no_network!(Ok(())); @@ -453,16 +615,15 @@ async fn turn_start_shell_zsh_fork_subcommand_decline_marks_parent_declined_v2() first_file.display(), second_file.display() ); - let tool_call_arguments = serde_json::to_string(&serde_json::json!({ - "command": shell_command, - "workdir": serde_json::Value::Null, - "timeout_ms": 5000 + let tool_call_arguments = serde_json::to_string(&json!({ + "cmd": shell_command, + "yield_time_ms": 5000, }))?; let response = responses::sse(vec![ responses::ev_response_created("resp-1"), responses::ev_function_call( "call-zsh-fork-subcommand-decline", - "shell_command", + "exec_command", &tool_call_arguments, ), responses::ev_completed("resp-1"), @@ -483,7 +644,7 @@ async fn turn_start_shell_zsh_fork_subcommand_decline_marks_parent_declined_v2() "untrusted", &BTreeMap::from([ (Feature::ShellZshFork, true), - (Feature::UnifiedExec, false), + (Feature::UnifiedExec, true), (Feature::ShellSnapshot, false), ]), &zsh_path, @@ -694,6 +855,21 @@ async fn create_zsh_test_mcp_process(codex_home: &Path, zdotdir: &Path) -> Resul McpProcess::new_with_env(codex_home, &[("ZDOTDIR", Some(zdotdir.as_str()))]).await } +fn create_zsh_fork_exec_command_sse_response( + command: &str, + call_id: &str, +) -> anyhow::Result { + let tool_call_arguments = serde_json::to_string(&json!({ + "cmd": command, + "yield_time_ms": 5000, + }))?; + Ok(responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call(call_id, "exec_command", &tool_call_arguments), + responses::ev_completed("resp-1"), + ])) +} + fn create_config_toml( codex_home: &Path, server_uri: &str, diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs index 1be2654cec..4b3d9db8e6 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -49,7 +49,6 @@ use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; -use tokio::sync::RwLock; use tokio_util::sync::CancellationToken; use uuid::Uuid; @@ -107,9 +106,6 @@ pub(super) async fn try_run_zsh_fork( req.timeout_ms .unwrap_or(crate::exec::DEFAULT_EXEC_COMMAND_TIMEOUT_MS), ); - let exec_policy = Arc::new(RwLock::new( - ctx.session.services.exec_policy.current().as_ref().clone(), - )); let command_executor = CoreShellCommandExecutor { command, cwd: sandbox_cwd, @@ -154,7 +150,6 @@ pub(super) async fn try_run_zsh_fork( let stopwatch = Stopwatch::new(effective_timeout); let cancel_token = stopwatch.cancellation_token(); let escalation_policy = CoreShellActionProvider { - policy: Arc::clone(&exec_policy), session: Arc::clone(&ctx.session), turn: Arc::clone(&ctx.turn), call_id: ctx.call_id.clone(), @@ -214,20 +209,30 @@ pub(crate) async fn prepare_unified_exec_zsh_fork( return Ok(None); } - let exec_policy = Arc::new(RwLock::new( - ctx.session.services.exec_policy.current().as_ref().clone(), - )); + let ExecRequest { + command, + cwd, + env, + network, + expiration: _expiration, + sandbox, + windows_sandbox_level, + sandbox_permissions, + sandbox_policy, + justification, + arg0, + } = &exec_request; let command_executor = CoreShellCommandExecutor { - command: exec_request.command.clone(), - cwd: exec_request.cwd.clone(), - sandbox_policy: exec_request.sandbox_policy.clone(), - sandbox: exec_request.sandbox, - env: exec_request.env.clone(), - network: exec_request.network.clone(), - windows_sandbox_level: exec_request.windows_sandbox_level, - sandbox_permissions: exec_request.sandbox_permissions, - justification: exec_request.justification.clone(), - arg0: exec_request.arg0.clone(), + command: command.clone(), + cwd: cwd.clone(), + sandbox_policy: sandbox_policy.clone(), + sandbox: *sandbox, + env: env.clone(), + network: network.clone(), + windows_sandbox_level: *windows_sandbox_level, + sandbox_permissions: *sandbox_permissions, + justification: justification.clone(), + arg0: arg0.clone(), sandbox_policy_cwd: ctx.turn.cwd.clone(), macos_seatbelt_profile_extensions: ctx .turn @@ -249,7 +254,6 @@ pub(crate) async fn prepare_unified_exec_zsh_fork( ) })?; let escalation_policy = CoreShellActionProvider { - policy: Arc::clone(&exec_policy), session: Arc::clone(&ctx.session), turn: Arc::clone(&ctx.turn), call_id: ctx.call_id.clone(), @@ -277,7 +281,6 @@ pub(crate) async fn prepare_unified_exec_zsh_fork( } struct CoreShellActionProvider { - policy: Arc>, session: Arc, turn: Arc, call_id: String, @@ -597,18 +600,16 @@ impl EscalationPolicy for CoreShellActionProvider { .await; } - let evaluation = { - let policy = self.policy.read().await; - evaluate_intercepted_exec_policy( - &policy, - program, - argv, - self.approval_policy, - &self.sandbox_policy, - self.sandbox_permissions, - ENABLE_INTERCEPTED_EXEC_POLICY_SHELL_WRAPPER_PARSING, - ) - }; + let policy = self.session.services.exec_policy.current(); + let evaluation = evaluate_intercepted_exec_policy( + policy.as_ref(), + program, + argv, + self.approval_policy, + &self.sandbox_policy, + self.sandbox_permissions, + ENABLE_INTERCEPTED_EXEC_POLICY_SHELL_WRAPPER_PARSING, + ); // When true, means the Evaluation was due to *.rules, not the // fallback function. let decision_driven_by_policy = diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index cabcf96cf9..09ef4c1ca3 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -118,8 +118,6 @@ impl ToolsConfig { let shell_type = if !features.enabled(Feature::ShellTool) { ConfigShellToolType::Disabled - } else if features.enabled(Feature::ShellZshFork) { - ConfigShellToolType::ShellCommand } else if features.enabled(Feature::UnifiedExec) { // If ConPTY not supported (for old Windows versions), fallback on ShellCommand. if codex_utils_pty::conpty_supported() { @@ -127,6 +125,8 @@ impl ToolsConfig { } else { ConfigShellToolType::ShellCommand } + } else if features.enabled(Feature::ShellZshFork) { + ConfigShellToolType::ShellCommand } else { model_info.shell_type }; @@ -2812,7 +2812,7 @@ mod tests { } #[test] - fn shell_zsh_fork_prefers_shell_command_over_unified_exec() { + fn shell_zsh_fork_uses_unified_exec_when_enabled() { let config = test_config(); let model_info = ModelsManager::construct_model_info_offline_for_tests("o3", &config); let mut features = Features::with_defaults(); @@ -2826,7 +2826,7 @@ mod tests { session_source: SessionSource::Cli, }); - assert_eq!(tools_config.shell_type, ConfigShellToolType::ShellCommand); + assert_eq!(tools_config.shell_type, ConfigShellToolType::UnifiedExec); assert_eq!( tools_config.shell_command_backend, ShellCommandBackendConfig::ZshFork diff --git a/codex-rs/core/tests/common/zsh_fork.rs b/codex-rs/core/tests/common/zsh_fork.rs index fbc46421d9..29db068ef5 100644 --- a/codex-rs/core/tests/common/zsh_fork.rs +++ b/codex-rs/core/tests/common/zsh_fork.rs @@ -91,6 +91,29 @@ where builder.build(server).await } +pub async fn build_unified_exec_zsh_fork_test( + server: &wiremock::MockServer, + runtime: ZshForkRuntime, + approval_policy: AskForApproval, + sandbox_policy: SandboxPolicy, + pre_build_hook: F, +) -> Result +where + F: FnOnce(&Path) + Send + 'static, +{ + let mut builder = test_codex() + .with_pre_build_hook(pre_build_hook) + .with_config(move |config| { + runtime.apply_to_config(config, approval_policy, sandbox_policy); + config.use_experimental_unified_exec_tool = true; + config + .features + .enable(Feature::UnifiedExec) + .expect("test config should allow feature update"); + }); + builder.build(server).await +} + fn find_test_zsh_path() -> Result> { let repo_root = codex_utils_cargo_bin::repo_root()?; let dotslash_zsh = repo_root.join("codex-rs/app-server/tests/suite/zsh"); diff --git a/codex-rs/core/tests/suite/approvals.rs b/codex-rs/core/tests/suite/approvals.rs index 9ca159965f..9bc8a9185b 100644 --- a/codex-rs/core/tests/suite/approvals.rs +++ b/codex-rs/core/tests/suite/approvals.rs @@ -35,6 +35,7 @@ use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use core_test_support::wait_for_event_with_timeout; +use core_test_support::zsh_fork::build_unified_exec_zsh_fork_test; use core_test_support::zsh_fork::build_zsh_fork_test; use core_test_support::zsh_fork::restrictive_workspace_write_policy; use core_test_support::zsh_fork::zsh_fork_runtime; @@ -1985,6 +1986,158 @@ async fn approving_execpolicy_amendment_persists_policy_and_skips_future_prompts Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[cfg(unix)] +async fn unified_exec_zsh_fork_execpolicy_amendment_skips_later_subcommands() -> Result<()> { + skip_if_no_network!(Ok(())); + + let Some(runtime) = zsh_fork_runtime("unified exec zsh-fork execpolicy amendment test")? else { + return Ok(()); + }; + + let approval_policy = AskForApproval::UnlessTrusted; + let sandbox_policy = SandboxPolicy::new_read_only_policy(); + let server = start_mock_server().await; + let test = build_unified_exec_zsh_fork_test( + &server, + runtime, + approval_policy, + sandbox_policy.clone(), + |_| {}, + ) + .await?; + let allow_prefix_path = test.cwd.path().join("allow-prefix-zsh-fork.txt"); + let _ = fs::remove_file(&allow_prefix_path); + + let call_id = "allow-prefix-zsh-fork"; + let command = "touch allow-prefix-zsh-fork.txt && touch allow-prefix-zsh-fork.txt"; + let event = exec_command_event( + call_id, + command, + Some(1_000), + SandboxPermissions::UseDefault, + None, + )?; + let _ = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-zsh-fork-allow-prefix-1"), + event, + ev_completed("resp-zsh-fork-allow-prefix-1"), + ]), + ) + .await; + let results = mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-zsh-fork-allow-prefix-1", "done"), + ev_completed("resp-zsh-fork-allow-prefix-2"), + ]), + ) + .await; + + submit_turn( + &test, + "allow-prefix-zsh-fork", + approval_policy, + sandbox_policy, + ) + .await?; + + let expected_execpolicy_amendment = ExecPolicyAmendment::new(vec![ + "touch".to_string(), + "allow-prefix-zsh-fork.txt".to_string(), + ]); + let mut saw_parent_approval = false; + let mut saw_subcommand_approval = false; + loop { + let event = wait_for_event(&test.codex, |event| { + matches!( + event, + EventMsg::ExecApprovalRequest(_) | EventMsg::TurnComplete(_) + ) + }) + .await; + + match event { + EventMsg::TurnComplete(_) => break, + EventMsg::ExecApprovalRequest(approval) => { + let command_parts = approval.command.clone(); + let last_arg = command_parts.last().map(String::as_str).unwrap_or_default(); + if last_arg == command { + assert!( + !saw_parent_approval, + "unexpected duplicate parent approval: {command_parts:?}" + ); + saw_parent_approval = true; + test.codex + .submit(Op::ExecApproval { + id: approval.effective_approval_id(), + turn_id: None, + decision: ReviewDecision::Approved, + }) + .await?; + continue; + } + + let is_touch_subcommand = command_parts + .iter() + .any(|part| part == "allow-prefix-zsh-fork.txt") + && command_parts + .first() + .is_some_and(|part| part.ends_with("/touch") || part == "touch"); + if is_touch_subcommand { + assert!( + !saw_subcommand_approval, + "execpolicy amendment should suppress later matching subcommand approvals: {command_parts:?}" + ); + saw_subcommand_approval = true; + assert_eq!( + approval.proposed_execpolicy_amendment, + Some(expected_execpolicy_amendment.clone()) + ); + test.codex + .submit(Op::ExecApproval { + id: approval.effective_approval_id(), + turn_id: None, + decision: ReviewDecision::ApprovedExecpolicyAmendment { + proposed_execpolicy_amendment: expected_execpolicy_amendment + .clone(), + }, + }) + .await?; + continue; + } + + test.codex + .submit(Op::ExecApproval { + id: approval.effective_approval_id(), + turn_id: None, + decision: ReviewDecision::Approved, + }) + .await?; + } + other => panic!("unexpected event: {other:?}"), + } + } + + assert!(saw_parent_approval, "expected parent unified-exec approval"); + assert!( + saw_subcommand_approval, + "expected at least one intercepted touch approval" + ); + + let result = parse_result(&results.single_request().function_call_output(call_id)); + assert_eq!(result.exit_code.unwrap_or(0), 0); + assert!( + allow_prefix_path.exists(), + "expected touch command to complete after approving the first intercepted subcommand; output: {}", + result.stdout + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[cfg(unix)] async fn matched_prefix_rule_runs_unsandboxed_under_zsh_fork() -> Result<()> { diff --git a/codex-rs/core/tests/suite/skill_approval.rs b/codex-rs/core/tests/suite/skill_approval.rs index d81cbb5c66..6822791093 100644 --- a/codex-rs/core/tests/suite/skill_approval.rs +++ b/codex-rs/core/tests/suite/skill_approval.rs @@ -18,6 +18,7 @@ use core_test_support::skip_if_no_network; use core_test_support::test_codex::TestCodex; use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; +use core_test_support::zsh_fork::build_unified_exec_zsh_fork_test; use core_test_support::zsh_fork::build_zsh_fork_test; use core_test_support::zsh_fork::restrictive_workspace_write_policy; use core_test_support::zsh_fork::zsh_fork_runtime; @@ -48,6 +49,13 @@ fn shell_command_arguments(command: &str) -> Result { }))?) } +fn exec_command_arguments(command: &str) -> Result { + Ok(serde_json::to_string(&json!({ + "cmd": command, + "yield_time_ms": 500, + }))?) +} + async fn submit_turn_with_policies( test: &TestCodex, prompt: &str, @@ -87,6 +95,38 @@ echo 'zsh-fork-stderr' >&2 ) } +#[cfg(unix)] +fn write_repo_skill_with_shell_script_contents( + repo_root: &Path, + name: &str, + script_name: &str, + script_contents: &str, +) -> Result { + use std::os::unix::fs::PermissionsExt; + + let skill_dir = repo_root.join(".agents").join("skills").join(name); + let scripts_dir = skill_dir.join("scripts"); + fs::create_dir_all(&scripts_dir)?; + fs::write(repo_root.join(".git"), "gitdir: here")?; + fs::write( + skill_dir.join("SKILL.md"), + format!( + r#"--- +name: {name} +description: {name} skill +--- +"# + ), + )?; + + let script_path = scripts_dir.join(script_name); + fs::write(&script_path, script_contents)?; + let mut permissions = fs::metadata(&script_path)?.permissions(); + permissions.set_mode(0o755); + fs::set_permissions(&script_path, permissions)?; + Ok(script_path) +} + #[cfg(unix)] fn write_skill_with_shell_script_contents( home: &Path, @@ -265,6 +305,168 @@ permissions: Ok(()) } +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unified_exec_zsh_fork_prompts_for_skill_script_execution() -> Result<()> { + skip_if_no_network!(Ok(())); + + let Some(runtime) = zsh_fork_runtime("unified exec zsh-fork skill prompt test")? else { + return Ok(()); + }; + + let server = start_mock_server().await; + let tool_call_id = "uexec-zsh-fork-skill-call"; + let test = build_unified_exec_zsh_fork_test( + &server, + runtime, + AskForApproval::OnRequest, + SandboxPolicy::new_workspace_write_policy(), + |home| { + write_skill_with_shell_script(home, "mbolin-test-skill", "hello-mbolin.sh").unwrap(); + write_skill_metadata( + home, + "mbolin-test-skill", + r#" +permissions: + file_system: + read: + - "./data" + write: + - "./output" +"#, + ) + .unwrap(); + }, + ) + .await?; + + let (script_path_str, command) = skill_script_command(&test, "hello-mbolin.sh")?; + let arguments = exec_command_arguments(&command)?; + let mocks = + mount_function_call_agent_response(&server, tool_call_id, &arguments, "exec_command").await; + + submit_turn_with_policies( + &test, + "use $mbolin-test-skill", + AskForApproval::OnRequest, + SandboxPolicy::new_workspace_write_policy(), + ) + .await?; + + let approval = wait_for_exec_approval_request(&test) + .await + .expect("expected exec approval request before completion"); + assert_eq!(approval.call_id, tool_call_id); + assert_eq!(approval.command, vec![script_path_str.clone()]); + assert_eq!( + approval.available_decisions, + Some(vec![ + ReviewDecision::Approved, + ReviewDecision::ApprovedForSession, + ReviewDecision::Abort, + ]) + ); + assert_eq!( + approval.additional_permissions, + Some(PermissionProfile { + file_system: Some(FileSystemPermissions { + read: Some(vec![absolute_path( + &test.codex_home_path().join("skills/mbolin-test-skill/data"), + )]), + write: Some(vec![absolute_path( + &test + .codex_home_path() + .join("skills/mbolin-test-skill/output"), + )]), + }), + ..Default::default() + }) + ); + + test.codex + .submit(Op::ExecApproval { + id: approval.effective_approval_id(), + turn_id: None, + decision: ReviewDecision::Denied, + }) + .await?; + + wait_for_turn_complete(&test).await; + + let call_output = mocks + .completion + .single_request() + .function_call_output(tool_call_id); + let output = call_output["output"].as_str().unwrap_or_default(); + assert!( + output.contains("Execution denied: User denied execution"), + "expected rejection marker in function_call_output: {output:?}" + ); + + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unified_exec_zsh_fork_keeps_skill_loading_pinned_to_turn_cwd() -> Result<()> { + skip_if_no_network!(Ok(())); + + let Some(runtime) = zsh_fork_runtime("unified exec zsh-fork turn cwd skill test")? else { + return Ok(()); + }; + + let server = start_mock_server().await; + let tool_call_id = "uexec-zsh-fork-repo-skill-call"; + let test = build_unified_exec_zsh_fork_test( + &server, + runtime, + AskForApproval::OnRequest, + SandboxPolicy::new_workspace_write_policy(), + |_| {}, + ) + .await?; + + let repo_root = test.cwd_path().join("repo"); + let script_path = write_repo_skill_with_shell_script_contents( + &repo_root, + "repo-skill", + "repo-skill.sh", + "#!/bin/sh\necho 'repo-skill-output'\n", + )?; + let script_path_quoted = shlex::try_join([script_path.to_string_lossy().as_ref()])?; + let repo_root_quoted = shlex::try_join([repo_root.to_string_lossy().as_ref()])?; + let command = format!("cd {repo_root_quoted} && {script_path_quoted}"); + let arguments = exec_command_arguments(&command)?; + let mocks = + mount_function_call_agent_response(&server, tool_call_id, &arguments, "exec_command").await; + + submit_turn_with_policies( + &test, + "run the repo skill after changing directories", + AskForApproval::OnRequest, + SandboxPolicy::new_workspace_write_policy(), + ) + .await?; + + let approval = wait_for_exec_approval_request(&test).await; + assert!( + approval.is_none(), + "changing directories inside unified exec should not load repo-local skills from the shell cwd", + ); + + let call_output = mocks + .completion + .single_request() + .function_call_output(tool_call_id); + let output = call_output["output"].as_str().unwrap_or_default(); + assert!( + output.contains("repo-skill-output"), + "expected repo skill script to run without skill-specific approval when only the shell cwd changes: {output:?}" + ); + + Ok(()) +} + /// Look for `additional_permissions == None`, then verify that both the first /// run and the cached session-approval rerun stay inside the turn sandbox. #[cfg(unix)]