mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
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
This commit is contained in:
@@ -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;
|
||||
|
||||
@@ -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")]
|
||||
|
||||
@@ -673,6 +673,7 @@ fn test_tool_runtime(session: Arc<Session>, turn_context: Arc<TurnContext>) -> 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(),
|
||||
|
||||
@@ -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::<crate::WaitForEnvironmentToolConfig>(),
|
||||
dynamic_tools: turn_context.dynamic_tools.as_slice(),
|
||||
},
|
||||
&sess.services.tool_search_handler_cache,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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<T>(arguments: &str) -> Result<T, FunctionCallError>
|
||||
where
|
||||
|
||||
@@ -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 `<environment_context>`.";
|
||||
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<ToolInvocation> for WaitForEnvironmentHandler {
|
||||
fn tool_name(&self) -> ToolName {
|
||||
@@ -34,17 +88,13 @@ impl ToolExecutor<ToolInvocation> 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 `<environment_context>`."
|
||||
.to_string(),
|
||||
)),
|
||||
JsonSchema::string(Some(self.environment_id_description.clone())),
|
||||
)]),
|
||||
/*required*/ Some(vec!["environment_id".to_string()]),
|
||||
/*additional_properties*/ Some(false.into()),
|
||||
|
||||
@@ -44,6 +44,7 @@ pub(crate) struct ToolRouterParams<'a> {
|
||||
pub(crate) tool_runtimes: Vec<Arc<dyn CoreToolRuntime>>,
|
||||
pub(crate) tool_suggest_candidates: Option<ToolSuggestCandidates>,
|
||||
pub(crate) extension_tool_executors: Vec<Arc<dyn ToolExecutor<ExtensionToolCall>>>,
|
||||
pub(crate) wait_for_environment_tool_config: Option<Arc<crate::WaitForEnvironmentToolConfig>>,
|
||||
pub(crate) dynamic_tools: &'a [DynamicToolSpec],
|
||||
}
|
||||
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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<dyn ToolExecutor<ExtensionToolCall>>],
|
||||
wait_for_environment_tool_config: Option<&'a Arc<crate::WaitForEnvironmentToolConfig>>,
|
||||
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 {
|
||||
|
||||
@@ -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<Arc<dyn CoreToolRuntime>>,
|
||||
tool_suggest_candidates: Option<ToolSuggestCandidates>,
|
||||
extension_tool_executors: Vec<Arc<dyn ToolExecutor<ExtensionToolCall>>>,
|
||||
wait_for_environment_tool_config: Option<Arc<WaitForEnvironmentToolConfig>>,
|
||||
dynamic_tools: Vec<DynamicToolSpec>,
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<Config> 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<TestCodex> {
|
||||
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("<environment_context>"))
|
||||
.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("<environment_context>"))
|
||||
.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;
|
||||
|
||||
Reference in New Issue
Block a user