From 81da9deb065d7adb283816b19b40f89bcc484276 Mon Sep 17 00:00:00 2001 From: TAFOYA-OAI Date: Fri, 24 Jul 2026 05:54:53 +0000 Subject: [PATCH] Allow hosts to customize `wait_for_environment` descriptions (#35106) ## What changed - Add `WaitForEnvironmentToolConfig` as thread extension data for overriding the model-visible tool and `environment_id` descriptions. - Preserve the default descriptions when no override is provided or when the configured descriptions exceed the input or serialized tool-spec limits. - Keep `wait_for_environment` availability gated by the deferred executor feature independently of whether a host override is present. ## Testing - Cover default, custom, oversized, and feature-disabled tool configurations. - Verify the custom descriptions in the deferred-environment integration flow. GitOrigin-RevId: 6b49a73a434becd99ea5df911be53f3706a17c0a --- codex-rs/core-api/src/lib.rs | 1 + codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/session/tests.rs | 2 + codex-rs/core/src/session/turn.rs | 4 + .../core/src/stream_events_utils_tests.rs | 1 + codex-rs/core/src/tools/handlers/mod.rs | 1 + .../tools/handlers/wait_for_environment.rs | 68 +++++++-- codex-rs/core/src/tools/router.rs | 1 + codex-rs/core/src/tools/router_tests.rs | 5 + codex-rs/core/src/tools/spec_plan.rs | 13 +- codex-rs/core/src/tools/spec_plan_tests.rs | 130 +++++++++++++++++ codex-rs/core/tests/suite/remote_env.rs | 132 +++++++++++++----- 12 files changed, 317 insertions(+), 42 deletions(-) diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index 3824f35589..de8031f871 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -35,6 +35,7 @@ pub use codex_core::StartThreadOptions; pub use codex_core::StateDbHandle; pub use codex_core::ThreadManager; pub use codex_core::ThreadShutdownReport; +pub use codex_core::WaitForEnvironmentToolConfig; pub use codex_core::build_models_manager; pub use codex_core::config::Config; pub use codex_core::config::Constrained; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index a3f29eb972..e37c17ca5d 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -116,6 +116,7 @@ pub use thread_manager::ThreadShutdownReport; pub use thread_manager::build_models_manager; pub use thread_manager::local_agent_graph_store_from_state_db; pub use thread_manager::thread_store_from_config; +pub use tools::handlers::WaitForEnvironmentToolConfig; pub use web_search::web_search_action_detail; pub use windows_sandbox_read_grants::grant_read_root_non_elevated; #[deprecated(note = "use ThreadManager")] diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 0b284dd5a7..308f39de98 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -673,6 +673,7 @@ fn test_tool_runtime(session: Arc, turn_context: Arc) -> T tool_suggest_candidates: None, tool_runtimes: Vec::new(), extension_tool_executors: Vec::new(), + wait_for_environment_tool_config: None, dynamic_tools: turn_context.dynamic_tools.as_slice(), }, &Default::default(), @@ -10563,6 +10564,7 @@ async fn fatal_tool_error_stops_turn_and_reports_error() { tool_suggest_candidates: None, tool_runtimes: Vec::new(), extension_tool_executors: Vec::new(), + wait_for_environment_tool_config: None, dynamic_tools: turn_context.dynamic_tools.as_slice(), }, &Default::default(), diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index a425affb19..71efc2aa7b 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1397,6 +1397,10 @@ pub(crate) async fn built_tools( tool_runtimes: mcp_tool_runtimes, tool_suggest_candidates, extension_tool_executors: extension_tool_executors(sess), + wait_for_environment_tool_config: sess + .services + .thread_extension_data + .get::(), dynamic_tools: turn_context.dynamic_tools.as_slice(), }, &sess.services.tool_search_handler_cache, diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index b481b7a840..85c4919978 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/codex-rs/core/src/stream_events_utils_tests.rs @@ -284,6 +284,7 @@ async fn handle_output_item_done_returns_contributed_last_agent_message() { tool_suggest_candidates: None, tool_runtimes: Vec::new(), extension_tool_executors: Vec::new(), + wait_for_environment_tool_config: None, dynamic_tools: turn_context.dynamic_tools.as_slice(), }, &Default::default(), diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index 3eaf39c3cf..3033e7b08c 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -78,6 +78,7 @@ pub(crate) use unified_exec::ExecCommandHandlerOptions; pub use unified_exec::WriteStdinHandler; pub use view_image::ViewImageHandler; pub(crate) use wait_for_environment::WaitForEnvironmentHandler; +pub use wait_for_environment::WaitForEnvironmentToolConfig; pub(crate) fn parse_arguments(arguments: &str) -> Result where diff --git a/codex-rs/core/src/tools/handlers/wait_for_environment.rs b/codex-rs/core/src/tools/handlers/wait_for_environment.rs index 6bcb6f3c62..b75ea60618 100644 --- a/codex-rs/core/src/tools/handlers/wait_for_environment.rs +++ b/codex-rs/core/src/tools/handlers/wait_for_environment.rs @@ -1,5 +1,3 @@ -use std::collections::BTreeMap; - use codex_tools::JsonSchema; use codex_tools::JsonToolOutput; use codex_tools::ResponsesApiTool; @@ -7,6 +5,7 @@ use codex_tools::ToolName; use codex_tools::ToolSpec; use serde::Deserialize; use serde_json::json; +use std::collections::BTreeMap; use crate::function_tool::FunctionCallError; use crate::tools::context::ToolInvocation; @@ -17,6 +16,24 @@ use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; const WAIT_FOR_ENVIRONMENT_TOOL_NAME: &str = "wait_for_environment"; +const DEFAULT_TOOL_DESCRIPTION: &str = "Wait for a selected execution environment marked as `starting` to become available. Use this when the current task needs that environment's files, commands, or installed capabilities. Do not wait if the task can be completed using tools already available, such as connectors. Waiting may take several minutes and blocks other tool calls. If startup fails, continue without that environment."; +const DEFAULT_ENVIRONMENT_ID_DESCRIPTION: &str = + "The exact environment ID marked as `starting` in ``."; +const MAX_COMBINED_DESCRIPTION_BYTES: usize = 1_024; +const MAX_SERIALIZED_TOOL_SPEC_BYTES: usize = 1_000; + +/// Model-visible descriptions supplied by a host that supports deferred environments. +/// +/// The two tool-schema descriptions must not exceed 1,024 UTF-8 bytes in total, and Core +/// also limits the complete serialized tool specification to 1,000 bytes. Oversized +/// descriptions fall back to Core defaults. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct WaitForEnvironmentToolConfig { + /// Explains when and why the model should call `wait_for_environment`. + pub tool_description: String, + /// Explains how the model should select the `environment_id` argument. + pub environment_id_description: String, +} #[derive(Deserialize)] #[serde(deny_unknown_fields)] @@ -24,7 +41,44 @@ struct WaitForEnvironmentArgs { environment_id: String, } -pub(crate) struct WaitForEnvironmentHandler; +pub(crate) struct WaitForEnvironmentHandler { + tool_description: String, + environment_id_description: String, +} + +impl WaitForEnvironmentHandler { + pub(crate) fn new(config: &WaitForEnvironmentToolConfig) -> Self { + let combined_description_bytes = config + .tool_description + .len() + .saturating_add(config.environment_id_description.len()); + if combined_description_bytes <= MAX_COMBINED_DESCRIPTION_BYTES { + let handler = Self { + tool_description: config.tool_description.clone(), + environment_id_description: config.environment_id_description.clone(), + }; + if serde_json::to_vec(&handler.spec()) + .is_ok_and(|serialized| serialized.len() <= MAX_SERIALIZED_TOOL_SPEC_BYTES) + { + return handler; + } + } + + tracing::warn!( + "oversized wait_for_environment tool configuration; falling back to Core defaults" + ); + Self::default() + } +} + +impl Default for WaitForEnvironmentHandler { + fn default() -> Self { + Self { + tool_description: DEFAULT_TOOL_DESCRIPTION.to_string(), + environment_id_description: DEFAULT_ENVIRONMENT_ID_DESCRIPTION.to_string(), + } + } +} impl ToolExecutor for WaitForEnvironmentHandler { fn tool_name(&self) -> ToolName { @@ -34,17 +88,13 @@ impl ToolExecutor for WaitForEnvironmentHandler { fn spec(&self) -> ToolSpec { ToolSpec::Function(ResponsesApiTool { name: WAIT_FOR_ENVIRONMENT_TOOL_NAME.to_string(), - description: "Wait for a selected execution environment marked as `starting` to become available. Use this when the current task needs that environment's files, commands, or installed capabilities. Do not wait if the task can be completed using tools already available, such as connectors. Waiting may take several minutes and blocks other tool calls. If startup fails, continue without that environment." - .to_string(), + description: self.tool_description.clone(), strict: false, defer_loading: None, parameters: JsonSchema::object( BTreeMap::from([( "environment_id".to_string(), - JsonSchema::string(Some( - "The exact environment ID marked as `starting` in ``." - .to_string(), - )), + JsonSchema::string(Some(self.environment_id_description.clone())), )]), /*required*/ Some(vec!["environment_id".to_string()]), /*additional_properties*/ Some(false.into()), diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 4dae7b3865..696beca2c1 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -44,6 +44,7 @@ pub(crate) struct ToolRouterParams<'a> { pub(crate) tool_runtimes: Vec>, pub(crate) tool_suggest_candidates: Option, pub(crate) extension_tool_executors: Vec>>, + pub(crate) wait_for_environment_tool_config: Option>, pub(crate) dynamic_tools: &'a [DynamicToolSpec], } diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index e0999e9120..9b8f90b22c 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -120,6 +120,7 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow tool_suggest_candidates: None, tool_runtimes: Vec::new(), extension_tool_executors: Vec::new(), + wait_for_environment_tool_config: None, dynamic_tools: turn.dynamic_tools.as_slice(), }, &Default::default(), @@ -233,6 +234,7 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { )), ], extension_tool_executors: Vec::new(), + wait_for_environment_tool_config: None, dynamic_tools: turn.dynamic_tools.as_slice(), }, &Default::default(), @@ -272,6 +274,7 @@ async fn tools_without_handlers_do_not_support_parallel() -> anyhow::Result<()> tool_suggest_candidates: None, tool_runtimes: Vec::new(), extension_tool_executors: Vec::new(), + wait_for_environment_tool_config: None, dynamic_tools: turn.dynamic_tools.as_slice(), }, &Default::default(), @@ -330,6 +333,7 @@ async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> { tool_suggest_candidates: None, tool_runtimes: Vec::new(), extension_tool_executors: Vec::new(), + wait_for_environment_tool_config: None, dynamic_tools: &dynamic_tools, }, &Default::default(), @@ -407,6 +411,7 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow tool_suggest_candidates: None, tool_runtimes: Vec::new(), extension_tool_executors: extension_tool_executors(&session), + wait_for_environment_tool_config: None, dynamic_tools: turn.dynamic_tools.as_slice(), }, &Default::default(), diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 6897f12b60..c9dbba0d33 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -147,6 +147,7 @@ struct CoreToolPlanContext<'a> { tool_runtimes: &'a [PlannedRuntime], tool_suggest_candidates: Option<&'a crate::tools::router::ToolSuggestCandidates>, extension_tool_executors: &'a [Arc>], + wait_for_environment_tool_config: Option<&'a Arc>, dynamic_tools: &'a [DynamicToolSpec], tool_search_handler_cache: &'a ToolSearchHandlerCache, default_agent_type_description: &'a str, @@ -183,6 +184,7 @@ fn build_tool_specs_and_registry( tool_runtimes, tool_suggest_candidates, extension_tool_executors, + wait_for_environment_tool_config, dynamic_tools, } = params; let default_agent_type_description = @@ -194,6 +196,7 @@ fn build_tool_specs_and_registry( tool_runtimes: &tool_runtimes, tool_suggest_candidates: tool_suggest_candidates.as_ref(), extension_tool_executors: &extension_tool_executors, + wait_for_environment_tool_config: wait_for_environment_tool_config.as_ref(), dynamic_tools, tool_search_handler_cache, default_agent_type_description: &default_agent_type_description, @@ -719,7 +722,15 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut } if features.enabled(Feature::DeferredExecutor) { - planned_tools.add(WaitForEnvironmentHandler); + planned_tools.add( + context + .wait_for_environment_tool_config + .map(Arc::as_ref) + .map_or_else( + WaitForEnvironmentHandler::default, + WaitForEnvironmentHandler::new, + ), + ); } if turn_context.config.experimental_request_user_input_enabled { diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 35526c5ab9..3957f42202 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -28,6 +28,7 @@ use codex_tools::ToolSpec; use pretty_assertions::assert_eq; use serde_json::json; +use crate::WaitForEnvironmentToolConfig; use crate::config::CurrentTimeReminderConfig; use crate::environment_selection::TurnEnvironmentState; use crate::session::step_context::StepContext; @@ -36,6 +37,7 @@ use crate::session::tests::mcp_config_for_test; use crate::session::turn_context::TurnContext; use crate::tools::handlers::McpHandler; use crate::tools::handlers::ToolSearchHandlerCache; +use crate::tools::handlers::WaitForEnvironmentHandler; use crate::tools::handlers::multi_agents_spec::MULTI_AGENT_V1_NAMESPACE; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::override_tool_exposure; @@ -51,6 +53,7 @@ struct ToolPlanInputs { tool_runtimes: Vec>, tool_suggest_candidates: Option, extension_tool_executors: Vec>>, + wait_for_environment_tool_config: Option>, dynamic_tools: Vec, } @@ -193,6 +196,7 @@ async fn probe_with( tool_suggest_candidates: inputs.tool_suggest_candidates, tool_runtimes: inputs.tool_runtimes, extension_tool_executors: inputs.extension_tool_executors, + wait_for_environment_tool_config: inputs.wait_for_environment_tool_config, dynamic_tools: inputs.dynamic_tools.as_slice(), }, &Default::default(), @@ -446,6 +450,129 @@ fn apply_patch_accepts_environment_id(spec: &ToolSpec) -> bool { } } +#[tokio::test] +async fn wait_for_environment_requires_feature_and_uses_host_config_when_present() { + const TOOL_DESCRIPTION: &str = "Host-provided wait tool description"; + const ENVIRONMENT_ID_DESCRIPTION: &str = "Host-provided environment ID description"; + + for deferred_executor_enabled in [false, true] { + for config_present in [false, true] { + let wait_for_environment_tool_config = config_present.then(|| { + Arc::new(WaitForEnvironmentToolConfig { + tool_description: TOOL_DESCRIPTION.to_string(), + environment_id_description: ENVIRONMENT_ID_DESCRIPTION.to_string(), + }) + }); + let plan = probe_with( + |turn| { + set_feature(turn, Feature::DeferredExecutor, deferred_executor_enabled); + }, + ToolPlanInputs { + wait_for_environment_tool_config, + ..ToolPlanInputs::default() + }, + ) + .await; + + if deferred_executor_enabled { + plan.assert_visible_contains(&["wait_for_environment"]); + plan.assert_registered_contains(&["wait_for_environment"]); + if !config_present { + assert_eq!( + plan.visible_spec("wait_for_environment"), + &WaitForEnvironmentHandler::default().spec() + ); + continue; + } + let ToolSpec::Function(ResponsesApiTool { + description, + parameters, + .. + }) = plan.visible_spec("wait_for_environment") + else { + panic!("expected wait_for_environment function spec"); + }; + assert_eq!(description, TOOL_DESCRIPTION); + assert_eq!( + parameters + .properties + .as_ref() + .and_then(|properties| properties.get("environment_id")) + .and_then(|schema| schema.description.as_deref()), + Some(ENVIRONMENT_ID_DESCRIPTION) + ); + } else { + plan.assert_visible_lacks(&["wait_for_environment"]); + plan.assert_registered_lacks(&["wait_for_environment"]); + } + } + } +} + +#[tokio::test] +async fn wait_for_environment_falls_back_for_oversized_host_configuration() { + const MAX_COMBINED_DESCRIPTION_BYTES: usize = 1_024; + + for (tool_description, environment_id_description) in [ + ( + "x".repeat(MAX_COMBINED_DESCRIPTION_BYTES + 1), + String::new(), + ), + ( + String::new(), + "x".repeat(MAX_COMBINED_DESCRIPTION_BYTES + 1), + ), + ("x".repeat(512), "x".repeat(513)), + // The descriptions fit the aggregate input cap, but the complete serialized schema does + // not fit its model-context cap once the surrounding tool definition is included. + ("x".repeat(500), "x".repeat(500)), + ] { + let configured_tool_description = tool_description.clone(); + let configured_environment_id_description = environment_id_description.clone(); + let plan = probe_with( + |turn| { + set_feature(turn, Feature::DeferredExecutor, /*enabled*/ true); + }, + ToolPlanInputs { + wait_for_environment_tool_config: Some(Arc::new(WaitForEnvironmentToolConfig { + tool_description, + environment_id_description, + })), + ..ToolPlanInputs::default() + }, + ) + .await; + + plan.assert_visible_contains(&["wait_for_environment"]); + plan.assert_registered_contains(&["wait_for_environment"]); + let ToolSpec::Function(ResponsesApiTool { + description, + parameters, + .. + }) = plan.visible_spec("wait_for_environment") + else { + panic!("expected wait_for_environment function spec"); + }; + let environment_id_description = parameters + .properties + .as_ref() + .and_then(|properties| properties.get("environment_id")) + .and_then(|schema| schema.description.as_deref()) + .expect("environment_id description should be present"); + assert_ne!(description, &configured_tool_description); + assert_ne!( + environment_id_description, + configured_environment_id_description + ); + assert!( + serde_json::to_vec(plan.visible_spec("wait_for_environment")) + .expect("tool spec should serialize") + .len() + <= 1_000 + ); + } +} + #[tokio::test] async fn request_user_input_tool_respects_experimental_config_gate() { let enabled = probe(|_| {}).await; @@ -710,6 +837,7 @@ async fn environment_tools_follow_the_step_context() { tool_runtimes: Vec::new(), tool_suggest_candidates: None, extension_tool_executors: Vec::new(), + wait_for_environment_tool_config: None, dynamic_tools: &[], }, &Default::default(), @@ -912,6 +1040,7 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { )], tool_suggest_candidates: None, extension_tool_executors: Vec::new(), + wait_for_environment_tool_config: None, dynamic_tools: &[], }, &cache, @@ -935,6 +1064,7 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { )], tool_suggest_candidates: None, extension_tool_executors: Vec::new(), + wait_for_environment_tool_config: None, dynamic_tools: &[], }, &cache, diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 05d4a38cff..ba07c6ca9e 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -3,7 +3,9 @@ use anyhow::Result; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_config::types::ApprovalsReviewer; +use codex_core::WaitForEnvironmentToolConfig; use codex_core::compact::SUMMARIZATION_PROMPT; +use codex_core::config::Config; use codex_core::config::Constrained; use codex_exec_server::CopyOptions; use codex_exec_server::CreateDirectoryOptions; @@ -16,6 +18,10 @@ use codex_exec_server::NoiseRendezvousConnectBundle; use codex_exec_server::NoiseRendezvousConnectProvider; use codex_exec_server::REMOTE_ENVIRONMENT_ID; use codex_exec_server::RemoveOptions; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadStartInput; use codex_features::Feature; use codex_protocol::models::FileSystemPermissions; use codex_protocol::models::PermissionProfile; @@ -66,6 +72,7 @@ use core_test_support::skip_if_no_remote_env; use core_test_support::skip_if_target_windows; use core_test_support::submit_thread_settings; use core_test_support::test_codex::TestCodex; +use core_test_support::test_codex::TestCodexBuilder; use core_test_support::test_codex::local; use core_test_support::test_codex::test_codex; use core_test_support::test_codex::test_env; @@ -97,6 +104,34 @@ use tokio::time::timeout; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::accept_async; use tokio_tungstenite::tungstenite::Message; + +const WAIT_FOR_ENVIRONMENT_TEST_TOOL_DESCRIPTION: &str = "Test wait tool description"; +const WAIT_FOR_ENVIRONMENT_TEST_ENVIRONMENT_ID_DESCRIPTION: &str = + "Test environment ID description"; + +struct WaitForEnvironmentTestExtension; + +impl ThreadLifecycleContributor for WaitForEnvironmentTestExtension { + fn on_thread_start<'a>( + &'a self, + input: ThreadStartInput<'a, Config>, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + input.thread_store.insert(WaitForEnvironmentToolConfig { + tool_description: WAIT_FOR_ENVIRONMENT_TEST_TOOL_DESCRIPTION.to_string(), + environment_id_description: WAIT_FOR_ENVIRONMENT_TEST_ENVIRONMENT_ID_DESCRIPTION + .to_string(), + }); + }) + } +} + +fn test_codex_with_wait_for_environment() -> TestCodexBuilder { + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.thread_lifecycle_contributor(Arc::new(WaitForEnvironmentTestExtension)); + test_codex().with_extensions(Arc::new(extensions.build())) +} + async fn unified_exec_test(server: &wiremock::MockServer) -> Result { let mut builder = test_codex().with_config(|config| { config.use_experimental_unified_exec_tool = true; @@ -338,36 +373,52 @@ async fn explicit_remote_shell_runs_in_remote_cwd() -> Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn step_world_state_does_not_duplicate_initial_environment_context() -> Result<()> { +async fn step_world_state_gates_deferred_prompt_independently_of_host_config() -> Result<()> { for deferred_executor_enabled in [false, true] { - let server = start_mock_server().await; - let response_mock = mount_sse_once( - &server, - sse(vec![ - ev_response_created("resp-1"), - ev_assistant_message("msg-1", "done"), - ev_completed("resp-1"), - ]), - ) - .await; - let mut builder = test_codex().with_config(move |config| { - if deferred_executor_enabled { - assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); - } - }); - let test = builder.build(&server).await?; + for host_config_present in [false, true] { + let server = start_mock_server().await; + let response_mock = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-1"), + ]), + ) + .await; + let builder = if host_config_present { + test_codex_with_wait_for_environment() + } else { + test_codex() + }; + let mut builder = builder.with_config(move |config| { + if deferred_executor_enabled { + assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); + } + }); + let test = builder.build(&server).await?; - test.submit_turn("report the environment").await?; + test.submit_turn("report the environment").await?; - let user_context = response_mock.single_request().message_input_texts("user"); - assert_eq!( - user_context - .iter() - .filter(|text| text.contains("")) - .count(), - 1, - "deferred executor enabled: {deferred_executor_enabled}", - ); + let request = response_mock.single_request(); + let user_context = request.message_input_texts("user"); + assert_eq!( + user_context + .iter() + .filter(|text| text.contains("")) + .count(), + 1, + "deferred executor enabled: {deferred_executor_enabled}; host config present: {host_config_present}", + ); + assert_eq!( + environment_instructions_occurrences(&request), + usize::from(deferred_executor_enabled), + ); + assert_eq!( + tool_names(&request.body_json()).contains(&"wait_for_environment".to_string()), + deferred_executor_enabled, + ); + } } Ok(()) @@ -679,7 +730,7 @@ async fn deferred_executor_starts_noise_connection_after_registration() -> Resul ], ) .await; - let mut builder = test_codex().with_config(|config| { + let mut builder = test_codex_with_wait_for_environment().with_config(|config| { config.use_experimental_unified_exec_tool = true; assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); assert!(config.features.enable(Feature::UnifiedExec).is_ok()); @@ -726,9 +777,26 @@ async fn deferred_executor_starts_noise_connection_after_registration() -> Resul let requests = response_mock.requests(); assert_eq!(requests.len(), 2); - let starting_tools = tool_names(&requests[0].body_json()); + let starting_request_body = requests[0].body_json(); + let starting_tools = tool_names(&starting_request_body); assert!(starting_tools.contains(&"wait_for_environment".to_string())); assert!(!starting_tools.contains(&"exec_command".to_string())); + let wait_tool = starting_request_body["tools"] + .as_array() + .and_then(|tools| { + tools + .iter() + .find(|tool| tool["name"] == "wait_for_environment") + }) + .context("wait_for_environment tool schema should be present")?; + assert_eq!( + wait_tool["description"].as_str(), + Some(WAIT_FOR_ENVIRONMENT_TEST_TOOL_DESCRIPTION) + ); + assert_eq!( + wait_tool["parameters"]["properties"]["environment_id"]["description"].as_str(), + Some(WAIT_FOR_ENVIRONMENT_TEST_ENVIRONMENT_ID_DESCRIPTION) + ); let (wait_output, _) = requests[1] .function_call_output_content_and_success(wait_call_id) .context("wait_for_environment output should be present")?; @@ -776,7 +844,7 @@ async fn deferred_executor_loads_agents_md_when_environment_becomes_ready() -> R ], ) .await; - let mut builder = test_codex() + let mut builder = test_codex_with_wait_for_environment() .with_exec_server_url(format!("ws://{}", listener.local_addr()?)) .with_config(|config| { assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); @@ -872,7 +940,7 @@ async fn deferred_executor_wait_reports_startup_failure() -> Result<()> { ], ) .await; - let mut builder = test_codex().with_config(|config| { + let mut builder = test_codex_with_wait_for_environment().with_config(|config| { config.use_experimental_unified_exec_tool = true; assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); assert!(config.features.enable(Feature::UnifiedExec).is_ok()); @@ -983,7 +1051,7 @@ async fn deferred_executor_compaction_preserves_then_updates_environment_once() ], ) .await; - let mut builder = test_codex() + let mut builder = test_codex_with_wait_for_environment() .with_exec_server_url(format!("ws://{}", listener.local_addr()?)) .with_config(|config| { config.project_doc_max_bytes = 0;