From fb0c5bb006de4c634b43c4d076289764c5d89b7e Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Tue, 23 Jun 2026 15:20:07 -0700 Subject: [PATCH] core: add wait tool for starting environments --- codex-rs/Cargo.lock | 12 ++ codex-rs/Cargo.toml | 2 + codex-rs/app-server/Cargo.toml | 1 + codex-rs/app-server/src/extensions.rs | 1 + codex-rs/core/Cargo.toml | 1 + codex-rs/core/src/environment_selection.rs | 18 ++ codex-rs/core/src/session/turn.rs | 2 +- codex-rs/core/src/tools/router.rs | 20 +- codex-rs/core/src/tools/router_tests.rs | 2 +- codex-rs/core/tests/suite/remote_env.rs | 203 +++++++++++++----- codex-rs/ext/deferred-executor/BUILD.bazel | 6 + codex-rs/ext/deferred-executor/Cargo.toml | 19 ++ codex-rs/ext/deferred-executor/src/lib.rs | 127 +++++++++++ .../src/capabilities/environment.rs | 27 +++ .../ext/extension-api/src/capabilities/mod.rs | 4 + .../ext/extension-api/src/contributors.rs | 24 ++- codex-rs/ext/extension-api/src/lib.rs | 4 + 17 files changed, 416 insertions(+), 57 deletions(-) create mode 100644 codex-rs/ext/deferred-executor/BUILD.bazel create mode 100644 codex-rs/ext/deferred-executor/Cargo.toml create mode 100644 codex-rs/ext/deferred-executor/src/lib.rs create mode 100644 codex-rs/ext/extension-api/src/capabilities/environment.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index a2953be530..3ffe1b3181 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -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" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 17c2d1561b..1fdc861dc5 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -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" } diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index 3d08059a63..23fa2f9fcd 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -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 } diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index d804e139bf..23b6128dfe 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -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); diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index d16e8d9cbe..988fb96912 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -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 } diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 1bed53fcfd..eb6be54e75 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -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, local_shell: Shell, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 100089723b..b13e4d9795 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -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, diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 4bdd121f6e..4c80683dad 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -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>> { + let starting_environments = step_context + .environments + .starting + .iter() + .cloned() + .map(|environment| Arc::new(environment) as Arc) + .collect::>(); 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() } diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index db1209cc03..9cd6b94edf 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -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(), diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index d360bde25f..4c870fab77 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -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> { + let mut builder = ExtensionRegistryBuilder::new(); + codex_deferred_executor_extension::install(&mut builder); + Arc::new(builder.build()) +} + +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_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::>() - }; - 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?; diff --git a/codex-rs/ext/deferred-executor/BUILD.bazel b/codex-rs/ext/deferred-executor/BUILD.bazel new file mode 100644 index 0000000000..9dfc496e70 --- /dev/null +++ b/codex-rs/ext/deferred-executor/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "deferred-executor", + crate_name = "codex_deferred_executor_extension", +) diff --git a/codex-rs/ext/deferred-executor/Cargo.toml b/codex-rs/ext/deferred-executor/Cargo.toml new file mode 100644 index 0000000000..b34097e50e --- /dev/null +++ b/codex-rs/ext/deferred-executor/Cargo.toml @@ -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 } diff --git a/codex-rs/ext/deferred-executor/src/lib.rs b/codex-rs/ext/deferred-executor/src/lib.rs new file mode 100644 index 0000000000..eaa71fccef --- /dev/null +++ b/codex-rs/ext/deferred-executor/src/lib.rs @@ -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>> { + Vec::new() + } + + fn tools_for_step( + &self, + input: ToolContributionInput<'_>, + ) -> Vec>> { + if input.starting_environments.is_empty() { + return Vec::new(); + } + + vec![Arc::new(WaitForEnvironmentTool { + environments: input.starting_environments.to_vec(), + })] + } +} + +struct WaitForEnvironmentTool { + environments: Vec>, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct WaitForEnvironmentArgs { + environment_id: String, +} + +impl ToolExecutor 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, 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(registry: &mut ExtensionRegistryBuilder) { + registry.tool_contributor(Arc::new(DeferredExecutorExtension)); +} diff --git a/codex-rs/ext/extension-api/src/capabilities/environment.rs b/codex-rs/ext/extension-api/src/capabilities/environment.rs new file mode 100644 index 0000000000..694cd3fb19 --- /dev/null +++ b/codex-rs/ext/extension-api/src/capabilities/environment.rs @@ -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 + 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<'_>; +} diff --git a/codex-rs/ext/extension-api/src/capabilities/mod.rs b/codex-rs/ext/extension-api/src/capabilities/mod.rs index 37c36e573b..25a2721c3f 100644 --- a/codex-rs/ext/extension-api/src/capabilities/mod.rs +++ b/codex-rs/ext/extension-api/src/capabilities/mod.rs @@ -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; diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 7d16efd3d4..cfacb4f7d6 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -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], +} + /// 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>>; + + /// Returns the complete native tool list for one model request. + fn tools_for_step( + &self, + input: ToolContributionInput<'_>, + ) -> Vec>> { + self.tools(input.session_store, input.thread_store) + } } /// Contributor for host-owned tool lifecycle gates. diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index de0e0c99f8..c3dab3bad2 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -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;