core: add wait tool for starting environments

This commit is contained in:
Sayan Sisodiya
2026-06-23 15:20:07 -07:00
parent bbbea91960
commit fb0c5bb006
17 changed files with 416 additions and 57 deletions
+12
View File
@@ -1977,6 +1977,7 @@ dependencies = [
"codex-config",
"codex-core",
"codex-core-plugins",
"codex-deferred-executor-extension",
"codex-exec-server",
"codex-extension-api",
"codex-external-agent-migration",
@@ -2626,6 +2627,7 @@ dependencies = [
"codex-context-fragments",
"codex-core-plugins",
"codex-core-skills",
"codex-deferred-executor-extension",
"codex-exec-server",
"codex-execpolicy",
"codex-extension-api",
@@ -2826,6 +2828,16 @@ dependencies = [
"zip 2.4.2",
]
[[package]]
name = "codex-deferred-executor-extension"
version = "0.0.0"
dependencies = [
"codex-extension-api",
"codex-tools",
"serde",
"serde_json",
]
[[package]]
name = "codex-exec"
version = "0.0.0"
+2
View File
@@ -47,6 +47,7 @@ members = [
"exec-server",
"execpolicy",
"execpolicy-legacy",
"ext/deferred-executor",
"ext/extension-api",
"ext/goal",
"ext/guardian",
@@ -169,6 +170,7 @@ codex-core = { path = "core" }
codex-core-api = { path = "core-api" }
codex-core-plugins = { path = "core-plugins" }
codex-core-skills = { path = "core-skills" }
codex-deferred-executor-extension = { path = "ext/deferred-executor" }
codex-exec = { path = "exec" }
codex-file-system = { path = "file-system" }
codex-exec-server = { path = "exec-server" }
+1
View File
@@ -35,6 +35,7 @@ codex-cloud-config = { workspace = true }
codex-config = { workspace = true }
codex-core = { workspace = true }
codex-core-plugins = { workspace = true }
codex-deferred-executor-extension = { workspace = true }
codex-home = { workspace = true }
codex-exec-server = { workspace = true }
codex-extension-api = { workspace = true }
+1
View File
@@ -71,6 +71,7 @@ where
|config: &Config| config.features.enabled(codex_features::Feature::Goals),
);
}
codex_deferred_executor_extension::install(&mut builder);
codex_guardian::install(&mut builder, guardian_agent_spawner);
codex_memories_extension::install(&mut builder, codex_otel::global());
codex_mcp_extension::install(&mut builder);
+1
View File
@@ -134,6 +134,7 @@ codex-shell-escalation = { workspace = true }
[dev-dependencies]
assert_cmd = { workspace = true }
assert_matches = { workspace = true }
codex-deferred-executor-extension = { workspace = true }
codex-image-generation-extension = { workspace = true }
codex-home = { workspace = true }
codex-otel = { workspace = true }
@@ -7,6 +7,9 @@ use codex_exec_server::Environment;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::ExecServerError;
use codex_exec_server::ExecutorFileSystem;
use codex_extension_api::EnvironmentStartupFuture;
use codex_extension_api::EnvironmentStartupOutcome;
use codex_extension_api::StartingEnvironment;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_path_uri::PathUri;
@@ -57,6 +60,21 @@ impl fmt::Debug for StartingTurnEnvironment {
}
}
impl StartingEnvironment for StartingTurnEnvironment {
fn environment_id(&self) -> &str {
&self.selection.environment_id
}
fn wait_until_ready(&self) -> EnvironmentStartupFuture<'_> {
Box::pin(async move {
match self.resolution.clone().await {
Ok(_) => EnvironmentStartupOutcome::Ready,
Err(_) => EnvironmentStartupOutcome::Failed,
}
})
}
}
pub(crate) struct ThreadEnvironments {
environment_manager: Arc<EnvironmentManager>,
local_shell: Shell,
+1 -1
View File
@@ -1362,7 +1362,7 @@ pub(crate) async fn built_tools(
mcp_tools,
deferred_mcp_tools,
tool_suggest_candidates,
extension_tool_executors: extension_tool_executors(sess),
extension_tool_executors: extension_tool_executors(sess, step_context),
dynamic_tools: turn_context.dynamic_tools.as_slice(),
},
&sess.services.tool_search_handler_cache,
+16 -4
View File
@@ -9,6 +9,8 @@ use crate::tools::registry::AnyToolResult;
use crate::tools::registry::ToolArgumentDiffConsumer;
use crate::tools::registry::ToolRegistry;
use crate::tools::spec_plan::build_tool_router;
use codex_extension_api::StartingEnvironment;
use codex_extension_api::ToolContributionInput;
use codex_mcp::ToolInfo;
use codex_protocol::dynamic_tools::DynamicToolSpec;
use codex_protocol::models::ResponseItem;
@@ -246,17 +248,27 @@ impl ToolRouter {
#[instrument(level = "trace", skip_all)]
pub(crate) fn extension_tool_executors(
session: &Session,
step_context: &StepContext,
) -> Vec<Arc<dyn ToolExecutor<ExtensionToolCall>>> {
let starting_environments = step_context
.environments
.starting
.iter()
.cloned()
.map(|environment| Arc::new(environment) as Arc<dyn StartingEnvironment>)
.collect::<Vec<_>>();
session
.services
.extensions
.tool_contributors()
.iter()
.flat_map(|contributor| {
contributor.tools(
&session.services.session_extension_data,
&session.services.thread_extension_data,
)
contributor.tools_for_step(ToolContributionInput {
session_store: &session.services.session_extension_data,
thread_store: &session.services.thread_extension_data,
turn_store: step_context.turn.extension_data.as_ref(),
starting_environments: &starting_environments,
})
})
.collect()
}
+1 -1
View File
@@ -365,7 +365,7 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow
tool_suggest_candidates: None,
deferred_mcp_tools: None,
mcp_tools: None,
extension_tool_executors: extension_tool_executors(&session),
extension_tool_executors: extension_tool_executors(&session, step_context.as_ref()),
dynamic_tools: turn.dynamic_tools.as_slice(),
},
&Default::default(),
+153 -50
View File
@@ -2,6 +2,7 @@ use anyhow::Context;
use anyhow::Result;
use codex_config::types::ApprovalsReviewer;
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;
@@ -9,9 +10,12 @@ use codex_exec_server::FileSystemSandboxContext;
use codex_exec_server::LOCAL_ENVIRONMENT_ID;
use codex_exec_server::REMOTE_ENVIRONMENT_ID;
use codex_exec_server::RemoveOptions;
use codex_extension_api::ExtensionRegistry;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_features::Feature;
use codex_protocol::models::FileSystemPermissions;
use codex_protocol::models::PermissionProfile;
use codex_protocol::models::ResponseItem;
use codex_protocol::permissions::FileSystemAccessMode;
use codex_protocol::permissions::FileSystemPath;
use codex_protocol::permissions::FileSystemSandboxEntry;
@@ -65,6 +69,7 @@ use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use std::process::Command;
use std::sync::Arc;
use std::time::Duration;
use std::time::SystemTime;
use std::time::UNIX_EPOCH;
@@ -401,32 +406,53 @@ async fn serve_environment_info(listener: TcpListener) {
.expect("environment info response");
}
fn deferred_executor_extensions() -> Arc<ExtensionRegistry<Config>> {
let mut builder = ExtensionRegistryBuilder::new();
codex_deferred_executor_extension::install(&mut builder);
Arc::new(builder.build())
}
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_raw_function_call(test: &TestCodex, call_id: &str) {
wait_for_event(&test.codex, |event| {
matches!(
event,
EventMsg::RawResponseItem(raw)
if matches!(
&raw.item,
ResponseItem::FunctionCall {
call_id: item_call_id,
..
} if item_call_id == call_id
)
)
})
.await;
}
#[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,
codex_deferred_executor_extension::WAIT_FOR_ENVIRONMENT_TOOL_NAME,
&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(),
),
@@ -456,6 +482,7 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<(
)
.await;
let mut builder = test_codex()
.with_extensions(deferred_executor_extensions())
.with_exec_server_url(format!("ws://{}", listener.local_addr()?))
.with_config(|config| {
config.use_experimental_unified_exec_tool = true;
@@ -469,12 +496,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 +513,8 @@ 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_raw_function_call(&test, wait_call_id).await;
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 +546,24 @@ 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(
&codex_deferred_executor_extension::WAIT_FOR_ENVIRONMENT_TOOL_NAME.to_string()
)
);
assert!(!starting_tools.contains(&"exec_command".to_string()));
assert!(ready_tools.contains(&"exec_command".to_string()));
assert!(
!ready_tools.contains(
&codex_deferred_executor_extension::WAIT_FOR_ENVIRONMENT_TOOL_NAME.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!(wait_output.is_some_and(|output| output.contains(r#""status":"ready""#)));
assert!(
requests[0]
.message_input_texts("user")
@@ -586,6 +597,98 @@ 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,
codex_deferred_executor_extension::WAIT_FOR_ENVIRONMENT_TOOL_NAME,
&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_extensions(deferred_executor_extensions())
.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_raw_function_call(&test, wait_call_id).await;
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(
&codex_deferred_executor_extension::WAIT_FOR_ENVIRONMENT_TOOL_NAME.to_string()
)
);
assert!(!starting_tools.contains(&"exec_command".to_string()));
assert!(
!failed_tools.contains(
&codex_deferred_executor_extension::WAIT_FOR_ENVIRONMENT_TOOL_NAME.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!(wait_output.is_some_and(|output| output.contains(r#""status":"failed""#)));
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?;
@@ -0,0 +1,6 @@
load("//:defs.bzl", "codex_rust_crate")
codex_rust_crate(
name = "deferred-executor",
crate_name = "codex_deferred_executor_extension",
)
+19
View File
@@ -0,0 +1,19 @@
[package]
edition.workspace = true
license.workspace = true
name = "codex-deferred-executor-extension"
version.workspace = true
[lib]
name = "codex_deferred_executor_extension"
path = "src/lib.rs"
doctest = false
[lints]
workspace = true
[dependencies]
codex-extension-api = { workspace = true }
codex-tools = { workspace = true }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
+127
View File
@@ -0,0 +1,127 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use codex_extension_api::EnvironmentStartupOutcome;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::FunctionCallError;
use codex_extension_api::JsonToolOutput;
use codex_extension_api::StartingEnvironment;
use codex_extension_api::ToolCall;
use codex_extension_api::ToolContributionInput;
use codex_extension_api::ToolContributor;
use codex_extension_api::ToolExecutor;
use codex_extension_api::ToolName;
use codex_extension_api::ToolOutput;
use codex_extension_api::ToolSpec;
use codex_tools::JsonSchema;
use codex_tools::ResponsesApiTool;
use serde::Deserialize;
use serde_json::json;
pub const WAIT_FOR_ENVIRONMENT_TOOL_NAME: &str = "wait_for_environment";
struct DeferredExecutorExtension;
impl ToolContributor for DeferredExecutorExtension {
fn tools(
&self,
_session_store: &ExtensionData,
_thread_store: &ExtensionData,
) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
Vec::new()
}
fn tools_for_step(
&self,
input: ToolContributionInput<'_>,
) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
if input.starting_environments.is_empty() {
return Vec::new();
}
vec![Arc::new(WaitForEnvironmentTool {
environments: input.starting_environments.to_vec(),
})]
}
}
struct WaitForEnvironmentTool {
environments: Vec<Arc<dyn StartingEnvironment>>,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct WaitForEnvironmentArgs {
environment_id: String,
}
impl ToolExecutor<ToolCall> for WaitForEnvironmentTool {
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()]),
Some(false.into()),
),
output_schema: None,
})
}
fn handle(&self, call: ToolCall) -> codex_extension_api::ToolExecutorFuture<'_> {
Box::pin(self.handle_call(call))
}
}
impl WaitForEnvironmentTool {
async fn handle_call(&self, call: ToolCall) -> Result<Box<dyn ToolOutput>, FunctionCallError> {
let args: WaitForEnvironmentArgs = serde_json::from_str(call.function_arguments()?)
.map_err(|err| FunctionCallError::RespondToModel(err.to_string()))?;
let Some(environment) = self
.environments
.iter()
.find(|environment| environment.environment_id() == args.environment_id)
else {
return Err(FunctionCallError::RespondToModel(format!(
"environment `{}` is not starting",
args.environment_id
)));
};
let environment_id = environment.environment_id();
let output = match environment.wait_until_ready().await {
EnvironmentStartupOutcome::Ready => JsonToolOutput::new(json!({
"environment_id": environment_id,
"status": "ready",
})),
EnvironmentStartupOutcome::Failed => JsonToolOutput::with_success(
json!({
"environment_id": environment_id,
"message": "The environment failed to start and is unavailable. Continue without it.",
"status": "failed",
}),
Some(false),
),
};
Ok(Box::new(output))
}
}
/// Installs the tool that waits for environments exposed as starting by the host.
pub fn install<C: Sync>(registry: &mut ExtensionRegistryBuilder<C>) {
registry.tool_contributor(Arc::new(DeferredExecutorExtension));
}
@@ -0,0 +1,27 @@
use std::future::Future;
use std::pin::Pin;
/// Future returned while an extension waits for a host environment to finish starting.
pub type EnvironmentStartupFuture<'a> =
Pin<Box<dyn Future<Output = EnvironmentStartupOutcome> + Send + 'a>>;
/// Host-observed result of waiting for an environment to finish starting.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EnvironmentStartupOutcome {
/// The environment is now available to the host.
Ready,
/// The environment failed to start and will not become available.
Failed,
}
/// Host-provided capability for an environment that is still starting.
///
/// Implementations must wait for the same startup operation that the host uses
/// to decide whether the environment is available to subsequent model requests.
pub trait StartingEnvironment: Send + Sync {
/// Stable host environment id used to select the environment.
fn environment_id(&self) -> &str;
/// Waits until the environment becomes available or permanently fails.
fn wait_until_ready(&self) -> EnvironmentStartupFuture<'_>;
}
@@ -1,9 +1,13 @@
mod agent;
mod environment;
mod events;
mod response_items;
pub use agent::AgentSpawnFuture;
pub use agent::AgentSpawner;
pub use environment::EnvironmentStartupFuture;
pub use environment::EnvironmentStartupOutcome;
pub use environment::StartingEnvironment;
pub use events::ExtensionEventSink;
pub use events::NoopExtensionEventSink;
pub use response_items::NoopResponseItemInjector;
+23 -1
View File
@@ -10,6 +10,7 @@ use codex_tools::ToolCall;
use codex_tools::ToolExecutor;
use crate::ExtensionData;
use crate::StartingEnvironment;
mod context;
mod mcp;
@@ -231,14 +232,35 @@ pub trait TokenUsageContributor: Send + Sync {
}
}
/// Host context available while extensions contribute tools for one model request.
#[derive(Clone, Copy)]
pub struct ToolContributionInput<'a> {
/// Store scoped to the host session runtime.
pub session_store: &'a ExtensionData,
/// Store scoped to this thread runtime.
pub thread_store: &'a ExtensionData,
/// Store scoped to this turn runtime.
pub turn_store: &'a ExtensionData,
/// Selected environments that are still starting for this request.
pub starting_environments: &'a [Arc<dyn StartingEnvironment>],
}
/// Extension contribution that exposes native tools owned by a feature.
pub trait ToolContributor: Send + Sync {
/// Returns the native tools visible for the supplied extension stores.
/// Returns tools that depend only on session- and thread-scoped state.
fn tools(
&self,
session_store: &ExtensionData,
thread_store: &ExtensionData,
) -> Vec<Arc<dyn ToolExecutor<ToolCall>>>;
/// Returns the complete native tool list for one model request.
fn tools_for_step(
&self,
input: ToolContributionInput<'_>,
) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
self.tools(input.session_store, input.thread_store)
}
}
/// Contributor for host-owned tool lifecycle gates.
+4
View File
@@ -6,11 +6,14 @@ mod user_instructions;
pub use capabilities::AgentSpawnFuture;
pub use capabilities::AgentSpawner;
pub use capabilities::EnvironmentStartupFuture;
pub use capabilities::EnvironmentStartupOutcome;
pub use capabilities::ExtensionEventSink;
pub use capabilities::NoopExtensionEventSink;
pub use capabilities::NoopResponseItemInjector;
pub use capabilities::ResponseItemInjectionFuture;
pub use capabilities::ResponseItemInjector;
pub use capabilities::StartingEnvironment;
pub use codex_context_fragments::ContextualUserFragment;
pub use codex_protocol::models::ResponseItem;
pub use codex_tools::ConversationHistory;
@@ -48,6 +51,7 @@ pub use contributors::ThreadStopInput;
pub use contributors::TokenUsageContributor;
pub use contributors::ToolCallOutcome;
pub use contributors::ToolCallSource;
pub use contributors::ToolContributionInput;
pub use contributors::ToolContributor;
pub use contributors::ToolFinishInput;
pub use contributors::ToolLifecycleContributor;