From 61ff4d087e23428335b2917e6def3fc47fab6425 Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Tue, 23 Jun 2026 17:35:34 -0700 Subject: [PATCH] core: add wait_for_environment for starting environments (#29745) ## Why With `DeferredExecutor`, a sampling request can begin while an environment is still starting. The model can see that pending state, but needs a way to wait for the environment within the same turn before continuing. Environment startup is owned by Core, so the wait tool should use the same request-frozen `StepContext` that advertised the starting environment. This keeps tool registration and execution tied to the exact startup operation the model saw, even if live thread state later changes. Supersedes #29735. ## What - register `wait_for_environment` when the current `StepContext` contains starting environments - wait on the selected `StartingTurnEnvironment` shared resolution and return a bounded ready or failed result - rebuild the next request normally, removing the wait tool and exposing ready environment tools, or reporting the environment as unavailable after failure ## Testing - `just test -p codex-core deferred_executor_` - verifies the wait tool is replaced by environment-backed tools after startup - verifies startup failure removes both the wait tool and unavailable environment tools while notifying the model --- codex-rs/core/src/environment_selection.rs | 6 + codex-rs/core/src/tools/handlers/mod.rs | 2 + .../tools/handlers/wait_for_environment.rs | 105 ++++++++++ codex-rs/core/src/tools/spec_plan.rs | 5 + codex-rs/core/tests/suite/remote_env.rs | 179 +++++++++++++----- 5 files changed, 247 insertions(+), 50 deletions(-) create mode 100644 codex-rs/core/src/tools/handlers/wait_for_environment.rs diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 1bed53fcf..fa86dded5 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/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> { + self.resolution.clone().await.map(|_| ()) + } +} + pub(crate) struct ThreadEnvironments { environment_manager: Arc, local_shell: Shell, diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index 47c81eabc..29ad17297 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/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(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 new file mode 100644 index 000000000..c9103521c --- /dev/null +++ b/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 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 {} diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 1ac255a83..bc4066cd2 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/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 { diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index d360bde25..e665e86a0 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/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 { + 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::>() - }; - 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::(&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?;