core: add wait_for_environment for starting environments

codex-rs/core/src/environment_selection.rs

@@ -57,6 +57,12 @@ impl fmt::Debug for StartingTurnEnvironment {
     }
 }
 
+impl StartingTurnEnvironment {
+    pub(crate) async fn wait_until_ready(&self) -> Result<(), Arc<ExecServerError>> {
+        self.resolution.clone().await.map(|_| ())
+    }
+}
+
 pub(crate) struct ThreadEnvironments {
     environment_manager: Arc<EnvironmentManager>,
     local_shell: Shell,

codex-rs/core/src/tools/handlers/mod.rs

@@ -35,6 +35,7 @@ pub(crate) mod tool_search_spec;
 pub(crate) mod unified_exec;
 mod view_image;
 pub(crate) mod view_image_spec;
+mod wait_for_environment;
 
 use codex_sandboxing::policy_transforms::intersect_permission_profiles;
 use codex_sandboxing::policy_transforms::merge_permission_profiles;
@@ -78,6 +79,7 @@ pub use unified_exec::ExecCommandHandler;
 pub(crate) use unified_exec::ExecCommandHandlerOptions;
 pub use unified_exec::WriteStdinHandler;
 pub use view_image::ViewImageHandler;
+pub(crate) use wait_for_environment::WaitForEnvironmentHandler;
 
 pub(crate) fn parse_arguments<T>(arguments: &str) -> Result<T, FunctionCallError>
 where

codex-rs/core/src/tools/handlers/wait_for_environment.rs

@@ -0,0 +1,105 @@
+use std::collections::BTreeMap;
+
+use codex_tools::JsonSchema;
+use codex_tools::JsonToolOutput;
+use codex_tools::ResponsesApiTool;
+use codex_tools::ToolName;
+use codex_tools::ToolSpec;
+use serde::Deserialize;
+use serde_json::json;
+
+use crate::function_tool::FunctionCallError;
+use crate::tools::context::ToolInvocation;
+use crate::tools::context::ToolPayload;
+use crate::tools::context::boxed_tool_output;
+use crate::tools::handlers::parse_arguments;
+use crate::tools::registry::CoreToolRuntime;
+use crate::tools::registry::ToolExecutor;
+
+const WAIT_FOR_ENVIRONMENT_TOOL_NAME: &str = "wait_for_environment";
+
+#[derive(Deserialize)]
+#[serde(deny_unknown_fields)]
+struct WaitForEnvironmentArgs {
+    environment_id: String,
+}
+
+pub(crate) struct WaitForEnvironmentHandler;
+
+impl ToolExecutor<ToolInvocation> for WaitForEnvironmentHandler {
+    fn tool_name(&self) -> ToolName {
+        ToolName::plain(WAIT_FOR_ENVIRONMENT_TOOL_NAME)
+    }
+
+    fn spec(&self) -> ToolSpec {
+        ToolSpec::Function(ResponsesApiTool {
+            name: WAIT_FOR_ENVIRONMENT_TOOL_NAME.to_string(),
+            description: "Wait for a starting environment to become available before continuing."
+                .to_string(),
+            strict: false,
+            defer_loading: None,
+            parameters: JsonSchema::object(
+                BTreeMap::from([(
+                    "environment_id".to_string(),
+                    JsonSchema::string(Some(
+                        "The id of an environment currently marked as starting.".to_string(),
+                    )),
+                )]),
+                /*required*/ Some(vec!["environment_id".to_string()]),
+                /*additional_properties*/ Some(false.into()),
+            ),
+            output_schema: None,
+        })
+    }
+
+    fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> {
+        Box::pin(async move {
+            let ToolInvocation {
+                payload,
+                step_context,
+                ..
+            } = invocation;
+            let arguments = match payload {
+                ToolPayload::Function { arguments } => arguments,
+                _ => {
+                    return Err(FunctionCallError::Fatal(format!(
+                        "{WAIT_FOR_ENVIRONMENT_TOOL_NAME} handler received unsupported payload"
+                    )));
+                }
+            };
+            let args: WaitForEnvironmentArgs = parse_arguments(&arguments)?;
+            let environment_id = args.environment_id;
+            let already_ready = step_context
+                .environments
+                .turn_environments
+                .iter()
+                .any(|environment| environment.environment_id == environment_id);
+            if !already_ready {
+                let Some(environment) = step_context
+                    .environments
+                    .starting
+                    .iter()
+                    .find(|environment| environment.selection.environment_id == environment_id)
+                    .cloned()
+                else {
+                    return Err(FunctionCallError::RespondToModel(format!(
+                        "environment `{environment_id}` is neither ready nor starting"
+                    )));
+                };
+
+                environment.wait_until_ready().await.map_err(|_| {
+                    FunctionCallError::RespondToModel(format!(
+                        "Environment `{environment_id}` failed to start and is unavailable. Continue without it."
+                    ))
+                })?;
+            }
+
+            Ok(boxed_tool_output(JsonToolOutput::new(json!({
+                "environment_id": environment_id,
+                "status": "ready",
+            }))))
+        })
+    }
+}
+
+impl CoreToolRuntime for WaitForEnvironmentHandler {}

codex-rs/core/src/tools/spec_plan.rs

@@ -29,6 +29,7 @@ use crate::tools::handlers::SleepHandler;
 use crate::tools::handlers::TestSyncHandler;
 use crate::tools::handlers::ToolSearchHandlerCache;
 use crate::tools::handlers::ViewImageHandler;
+use crate::tools::handlers::WaitForEnvironmentHandler;
 use crate::tools::handlers::WriteStdinHandler;
 use crate::tools::handlers::agent_jobs::ReportAgentJobResultHandler;
 use crate::tools::handlers::agent_jobs::SpawnAgentsOnCsvHandler;
@@ -719,6 +720,10 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut
 
     planned_tools.add(PlanHandler);
 
+    if features.enabled(Feature::DeferredExecutor) {
+        planned_tools.add(WaitForEnvironmentHandler);
+    }
+
     if turn_context.config.experimental_request_user_input_enabled {
         planned_tools.add_with_exposure(
             RequestUserInputHandler {

codex-rs/core/tests/suite/remote_env.rs

@@ -35,6 +35,7 @@ use codex_utils_path_uri::PathUri;
 use core_test_support::PathBufExt;
 use core_test_support::PathExt;
 use core_test_support::TestTargetOs;
+use core_test_support::responses::ResponseMock;
 use core_test_support::responses::ev_apply_patch_custom_tool_call;
 use core_test_support::responses::ev_assistant_message;
 use core_test_support::responses::ev_completed;
@@ -401,32 +402,40 @@ async fn serve_environment_info(listener: TcpListener) {
         .expect("environment info response");
 }
 
+fn tool_names(body: &Value) -> Vec<String> {
+    body["tools"]
+        .as_array()
+        .expect("tools array")
+        .iter()
+        .filter_map(|tool| tool.get("name").and_then(Value::as_str).map(str::to_owned))
+        .collect()
+}
+
+async fn wait_for_response_request_count(response_mock: &ResponseMock, expected_count: usize) {
+    timeout(Duration::from_secs(5), async {
+        while response_mock.requests().len() < expected_count {
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
+    })
+    .await
+    .expect("timed out waiting for Responses API request");
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<()> {
     let listener = TcpListener::bind("127.0.0.1:0").await?;
     let server = start_mock_server().await;
-    let user_input_call_id = "wait-for-startup";
+    let wait_call_id = "wait-for-startup";
     let response_mock = mount_sse_sequence(
         &server,
         vec![
             sse(vec![
                 ev_response_created("resp-1"),
                 ev_function_call(
-                    user_input_call_id,
-                    "request_user_input",
+                    wait_call_id,
+                    "wait_for_environment",
                     &json!({
-                        "questions": [{
-                            "id": "continue",
-                            "header": "Continue",
-                            "question": "Continue after startup?",
-                            "options": [{
-                                "label": "Yes (Recommended)",
-                                "description": "Continue the test."
-                            }, {
-                                "label": "No",
-                                "description": "Stop the test."
-                            }]
-                        }]
+                        "environment_id": REMOTE_ENVIRONMENT_ID,
                     })
                     .to_string(),
                 ),
@@ -469,12 +478,6 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<(
                     .enable(Feature::RequestPermissionsTool)
                     .is_ok()
             );
-            assert!(
-                config
-                    .features
-                    .enable(Feature::DefaultModeRequestUserInput)
-                    .is_ok()
-            );
         });
     let test = timeout(Duration::from_secs(5), builder.build(&server))
         .await
@@ -492,26 +495,9 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<(
             thread_settings: Default::default(),
         })
         .await?;
-    let request = wait_for_event_match(&test.codex, |event| match event {
-        EventMsg::RequestUserInput(request) => Some(request.clone()),
-        _ => None,
-    })
-    .await;
-
+    wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await;
+    assert_eq!(response_mock.requests().len(), 1);
     serve_environment_info(listener).await;
-    test.codex
-        .submit(Op::UserInputAnswer {
-            id: request.turn_id,
-            response: RequestUserInputResponse {
-                answers: HashMap::from([(
-                    "continue".to_string(),
-                    RequestUserInputAnswer {
-                        answers: vec!["Yes (Recommended)".to_string()],
-                    },
-                )]),
-            },
-        })
-        .await?;
     let event = wait_for_event(&test.codex, |event| {
         matches!(
             event,
@@ -543,16 +529,22 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<(
 
     let requests = response_mock.requests();
     assert_eq!(requests.len(), 3);
-    let tool_names = |request_index: usize| {
-        requests[request_index].body_json()["tools"]
-            .as_array()
-            .expect("tools array")
-            .iter()
-            .filter_map(|tool| tool.get("name").and_then(Value::as_str).map(str::to_owned))
-            .collect::<Vec<_>>()
-    };
-    assert!(!tool_names(0).contains(&"exec_command".to_string()));
-    assert!(tool_names(1).contains(&"exec_command".to_string()));
+    let starting_tools = tool_names(&requests[0].body_json());
+    let ready_tools = tool_names(&requests[1].body_json());
+    assert!(starting_tools.contains(&"wait_for_environment".to_string()));
+    assert!(!starting_tools.contains(&"exec_command".to_string()));
+    assert!(ready_tools.contains(&"exec_command".to_string()));
+    assert!(ready_tools.contains(&"wait_for_environment".to_string()));
+    let (wait_output, _) = requests[1]
+        .function_call_output_content_and_success(wait_call_id)
+        .context("wait_for_environment output should be present")?;
+    assert_eq!(
+        serde_json::from_str::<Value>(&wait_output.context("wait output should contain text")?)?,
+        json!({
+            "environment_id": REMOTE_ENVIRONMENT_ID,
+            "status": "ready",
+        })
+    );
     assert!(
         requests[0]
             .message_input_texts("user")
@@ -586,6 +578,93 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<(
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn deferred_executor_wait_reports_startup_failure() -> Result<()> {
+    let listener = TcpListener::bind("127.0.0.1:0").await?;
+    let server = start_mock_server().await;
+    let wait_call_id = "wait-for-failure";
+    let response_mock = mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-1"),
+                ev_function_call(
+                    wait_call_id,
+                    "wait_for_environment",
+                    &json!({
+                        "environment_id": REMOTE_ENVIRONMENT_ID,
+                    })
+                    .to_string(),
+                ),
+                ev_completed("resp-1"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-2"),
+                ev_assistant_message("msg-2", "done"),
+                ev_completed("resp-2"),
+            ]),
+        ],
+    )
+    .await;
+    let mut builder = test_codex()
+        .with_exec_server_url(format!("ws://{}", listener.local_addr()?))
+        .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());
+        });
+    let test = timeout(Duration::from_secs(5), builder.build(&server))
+        .await
+        .context("thread startup should not wait for the remote environment")??;
+
+    test.codex
+        .submit(Op::UserInput {
+            items: vec![UserInput::Text {
+                text: "wait for the environment".into(),
+                text_elements: Vec::new(),
+            }],
+            final_output_json_schema: None,
+            responsesapi_client_metadata: None,
+            additional_context: Default::default(),
+            thread_settings: Default::default(),
+        })
+        .await?;
+    wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await;
+    assert_eq!(response_mock.requests().len(), 1);
+    let (stream, _) = timeout(Duration::from_secs(5), listener.accept())
+        .await
+        .context("exec-server connection should arrive")??;
+    drop(stream);
+    wait_for_event(&test.codex, |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+
+    let requests = response_mock.requests();
+    assert_eq!(requests.len(), 2);
+    let starting_tools = tool_names(&requests[0].body_json());
+    let failed_tools = tool_names(&requests[1].body_json());
+    assert!(starting_tools.contains(&"wait_for_environment".to_string()));
+    assert!(!starting_tools.contains(&"exec_command".to_string()));
+    assert!(failed_tools.contains(&"wait_for_environment".to_string()));
+    assert!(!failed_tools.contains(&"exec_command".to_string()));
+    let (wait_output, _) = requests[1]
+        .function_call_output_content_and_success(wait_call_id)
+        .context("wait_for_environment output should be present")?;
+    assert_eq!(
+        wait_output.as_deref(),
+        Some("Environment `remote` failed to start and is unavailable. Continue without it.")
+    );
+    assert!(
+        requests[1]
+            .message_input_texts("user")
+            .iter()
+            .any(|text| text.contains("status=\"unavailable\""))
+    );
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn deferred_executor_compaction_preserves_then_updates_environment_once() -> Result<()> {
     let listener = TcpListener::bind("127.0.0.1:0").await?;