mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
Support pending remote environment registration (#32231)
## Why Remote environment provisioning can finish after a thread starts, before an exec-server WebSocket URL is available. ## What changed - Add `EnvironmentManager::register_pending_environment` and a one-shot `PendingEnvironmentRegistration` handle that resolves to either a validated WebSocket URL or a terminal provisioning error. - Let lazy remote exec-server clients wait for that result, while preserving reconnection behavior after a successful registration. - Keep replacement registrations isolated so completing an older handle does not resolve the current environment with the same ID. ## Testing Add coverage for successful connection and reconnection, provisioning and dropped-registration failures, invalid URLs, replacement isolation, and the deferred-executor startup flow. GitOrigin-RevId: 5c05be2b72291b77a1f71176d7075b1ad63332a5
This commit is contained in:
@@ -1843,7 +1843,7 @@ impl TestAppServerBuilder {
|
||||
let mut auto_env_overrides = vec![
|
||||
(
|
||||
CODEX_EXEC_SERVER_URL_ENV_VAR.to_string(),
|
||||
auto_env.environment().exec_server_url().map(str::to_string),
|
||||
auto_env.exec_server_url().map(str::to_string),
|
||||
),
|
||||
(
|
||||
CODEX_EXEC_SERVER_NOISE_REGISTRY_URL_ENV_VAR.to_string(),
|
||||
|
||||
@@ -163,6 +163,10 @@ impl TestEnv {
|
||||
&self.environment
|
||||
}
|
||||
|
||||
pub fn exec_server_url(&self) -> Option<&str> {
|
||||
self.exec_server_url.as_deref()
|
||||
}
|
||||
|
||||
/// Returns the environment and target-native cwd selected by the test harness.
|
||||
pub fn selection(&self) -> &TurnEnvironmentSelection {
|
||||
&self.selection
|
||||
|
||||
@@ -27,7 +27,9 @@ use codex_protocol::protocol::ReviewDecision;
|
||||
use codex_protocol::protocol::RolloutItem;
|
||||
use codex_protocol::protocol::RolloutLine;
|
||||
use codex_protocol::protocol::SandboxPolicy;
|
||||
use codex_protocol::protocol::ThreadSettingsOverrides;
|
||||
use codex_protocol::protocol::TurnEnvironmentSelection;
|
||||
use codex_protocol::protocol::TurnEnvironmentSelections;
|
||||
use codex_protocol::request_permissions::PermissionGrantScope;
|
||||
use codex_protocol::request_permissions::RequestPermissionProfile;
|
||||
use codex_protocol::request_permissions::RequestPermissionsResponse;
|
||||
@@ -530,25 +532,26 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<(
|
||||
],
|
||||
)
|
||||
.await;
|
||||
let mut builder = test_codex()
|
||||
.with_exec_server_url(format!("ws://{}", listener.local_addr()?))
|
||||
.with_config(|config| {
|
||||
config.project_doc_max_bytes = 0;
|
||||
config.use_experimental_unified_exec_tool = true;
|
||||
config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
|
||||
config.approvals_reviewer = ApprovalsReviewer::User;
|
||||
assert!(config.features.enable(Feature::DeferredExecutor).is_ok());
|
||||
assert!(config.features.enable(Feature::UnifiedExec).is_ok());
|
||||
assert!(
|
||||
config
|
||||
.features
|
||||
.enable(Feature::RequestPermissionsTool)
|
||||
.is_ok()
|
||||
);
|
||||
});
|
||||
let mut builder = test_codex().with_config(|config| {
|
||||
config.project_doc_max_bytes = 0;
|
||||
config.use_experimental_unified_exec_tool = true;
|
||||
config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
|
||||
config.approvals_reviewer = ApprovalsReviewer::User;
|
||||
assert!(config.features.enable(Feature::DeferredExecutor).is_ok());
|
||||
assert!(config.features.enable(Feature::UnifiedExec).is_ok());
|
||||
assert!(
|
||||
config
|
||||
.features
|
||||
.enable(Feature::RequestPermissionsTool)
|
||||
.is_ok()
|
||||
);
|
||||
});
|
||||
let test = timeout(Duration::from_secs(5), builder.build(&server))
|
||||
.await
|
||||
.context("thread startup should not wait for the remote environment")??;
|
||||
let environment_manager = test.thread_manager.environment_manager();
|
||||
let registration =
|
||||
environment_manager.register_pending_environment(REMOTE_ENVIRONMENT_ID.to_string())?;
|
||||
|
||||
test.codex
|
||||
.submit(Op::UserInput {
|
||||
@@ -559,11 +562,21 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<(
|
||||
final_output_json_schema: None,
|
||||
responsesapi_client_metadata: None,
|
||||
additional_context: Default::default(),
|
||||
thread_settings: Default::default(),
|
||||
thread_settings: ThreadSettingsOverrides {
|
||||
environments: Some(TurnEnvironmentSelections::new(
|
||||
test.config.cwd.clone(),
|
||||
vec![TurnEnvironmentSelection {
|
||||
environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
|
||||
cwd: PathUri::from_abs_path(&test.config.cwd),
|
||||
}],
|
||||
)),
|
||||
..Default::default()
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await;
|
||||
assert_eq!(response_mock.requests().len(), 1);
|
||||
registration.complete(Ok(format!("ws://{}", listener.local_addr()?)))?;
|
||||
serve_environment_info(listener).await;
|
||||
let event = wait_for_event(&test.codex, |event| {
|
||||
matches!(
|
||||
@@ -740,7 +753,6 @@ fn agents_md_occurrences(request: &ResponsesRequest, contents: &str) -> usize {
|
||||
|
||||
#[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(
|
||||
@@ -766,16 +778,17 @@ async fn deferred_executor_wait_reports_startup_failure() -> Result<()> {
|
||||
],
|
||||
)
|
||||
.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 mut builder = test_codex().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")??;
|
||||
let environment_manager = test.thread_manager.environment_manager();
|
||||
let registration =
|
||||
environment_manager.register_pending_environment(REMOTE_ENVIRONMENT_ID.to_string())?;
|
||||
|
||||
test.codex
|
||||
.submit(Op::UserInput {
|
||||
@@ -786,15 +799,21 @@ async fn deferred_executor_wait_reports_startup_failure() -> Result<()> {
|
||||
final_output_json_schema: None,
|
||||
responsesapi_client_metadata: None,
|
||||
additional_context: Default::default(),
|
||||
thread_settings: Default::default(),
|
||||
thread_settings: ThreadSettingsOverrides {
|
||||
environments: Some(TurnEnvironmentSelections::new(
|
||||
test.config.cwd.clone(),
|
||||
vec![TurnEnvironmentSelection {
|
||||
environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
|
||||
cwd: PathUri::from_abs_path(&test.config.cwd),
|
||||
}],
|
||||
)),
|
||||
..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);
|
||||
registration.complete(Err("CCA provisioning failed".to_string()))?;
|
||||
wait_for_event(&test.codex, |event| {
|
||||
matches!(event, EventMsg::TurnComplete(_))
|
||||
})
|
||||
|
||||
@@ -427,6 +427,7 @@ impl LazyRemoteExecServerClient {
|
||||
matches!(
|
||||
self.transport_params,
|
||||
ExecServerTransportParams::WebSocketUrl { .. }
|
||||
| ExecServerTransportParams::PendingWebSocketUrl(..)
|
||||
| ExecServerTransportParams::NoiseRendezvous { .. }
|
||||
)
|
||||
}
|
||||
|
||||
@@ -4,6 +4,8 @@ use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
use futures::future::Shared;
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
use crate::ExecServerError;
|
||||
use crate::HttpRequestParams;
|
||||
@@ -88,6 +90,8 @@ pub(crate) struct StdioExecServerCommand {
|
||||
pub cwd: Option<PathBuf>,
|
||||
}
|
||||
|
||||
pub(crate) type PendingExecServerUrl = Shared<oneshot::Receiver<Result<String, String>>>;
|
||||
|
||||
/// Parameters used to connect to a remote exec-server environment.
|
||||
#[derive(Clone)]
|
||||
pub(crate) enum ExecServerTransportParams {
|
||||
@@ -96,6 +100,7 @@ pub(crate) enum ExecServerTransportParams {
|
||||
connect_timeout: Duration,
|
||||
initialize_timeout: Duration,
|
||||
},
|
||||
PendingWebSocketUrl(PendingExecServerUrl),
|
||||
NoiseRendezvous {
|
||||
provider: Arc<dyn NoiseRendezvousConnectProvider>,
|
||||
identity: NoiseChannelIdentity,
|
||||
@@ -120,6 +125,9 @@ impl std::fmt::Debug for ExecServerTransportParams {
|
||||
.field("connect_timeout", connect_timeout)
|
||||
.field("initialize_timeout", initialize_timeout)
|
||||
.finish(),
|
||||
Self::PendingWebSocketUrl(..) => {
|
||||
f.debug_tuple("PendingWebSocketUrl").finish_non_exhaustive()
|
||||
}
|
||||
Self::NoiseRendezvous { .. } => {
|
||||
f.debug_struct("NoiseRendezvous").finish_non_exhaustive()
|
||||
}
|
||||
|
||||
@@ -19,6 +19,7 @@ use crate::ExecServerError;
|
||||
use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
|
||||
use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT;
|
||||
use crate::client_api::ExecServerClientConnectOptions;
|
||||
use crate::client_api::ExecServerTransportParams;
|
||||
use crate::client_api::NoiseRendezvousConnectArgs;
|
||||
use crate::client_api::NoiseRendezvousConnectBundle;
|
||||
use crate::client_api::NoiseRendezvousConnectProvider;
|
||||
@@ -91,27 +92,30 @@ impl ExecServerClient {
|
||||
/// Noise connection details are fetched here so reconnects get a fresh URL
|
||||
/// and authorization without replacing the harness identity.
|
||||
pub(crate) async fn connect_for_transport(
|
||||
transport_params: crate::client_api::ExecServerTransportParams,
|
||||
transport_params: ExecServerTransportParams,
|
||||
) -> Result<Self, ExecServerError> {
|
||||
match transport_params {
|
||||
crate::client_api::ExecServerTransportParams::WebSocketUrl {
|
||||
let (websocket_url, connect_timeout, initialize_timeout) = match transport_params {
|
||||
ExecServerTransportParams::WebSocketUrl {
|
||||
websocket_url,
|
||||
connect_timeout,
|
||||
initialize_timeout,
|
||||
} => {
|
||||
Self::connect_websocket(RemoteExecServerConnectArgs {
|
||||
} => (websocket_url, connect_timeout, initialize_timeout),
|
||||
ExecServerTransportParams::PendingWebSocketUrl(websocket_url) => {
|
||||
let websocket_url = websocket_url
|
||||
.await
|
||||
.unwrap_or_else(|_| {
|
||||
Err("environment registration ended before completion".to_string())
|
||||
})
|
||||
.map_err(|message| {
|
||||
ExecServerError::Disconnected(format!("environment unavailable: {message}"))
|
||||
})?;
|
||||
(
|
||||
websocket_url,
|
||||
client_name: ENVIRONMENT_CLIENT_NAME.to_string(),
|
||||
connect_timeout,
|
||||
initialize_timeout,
|
||||
resume_session_id: None,
|
||||
})
|
||||
.await
|
||||
DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
|
||||
DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
|
||||
)
|
||||
}
|
||||
crate::client_api::ExecServerTransportParams::NoiseRendezvous {
|
||||
provider,
|
||||
identity,
|
||||
} => {
|
||||
ExecServerTransportParams::NoiseRendezvous { provider, identity } => {
|
||||
let reconnect_strategy = ExecServerReconnectStrategy::NoiseRendezvous {
|
||||
provider: Arc::clone(&provider),
|
||||
identity: identity.clone(),
|
||||
@@ -121,21 +125,30 @@ impl ExecServerClient {
|
||||
};
|
||||
let (connection, options) =
|
||||
Self::open_initial_noise_rendezvous_connection(&provider, &identity).await?;
|
||||
Self::connect_with_recovery(connection, options, Some(reconnect_strategy)).await
|
||||
return Self::connect_with_recovery(connection, options, Some(reconnect_strategy))
|
||||
.await;
|
||||
}
|
||||
crate::client_api::ExecServerTransportParams::StdioCommand {
|
||||
ExecServerTransportParams::StdioCommand {
|
||||
command,
|
||||
initialize_timeout,
|
||||
} => {
|
||||
Self::connect_stdio_command(StdioExecServerConnectArgs {
|
||||
return Self::connect_stdio_command(StdioExecServerConnectArgs {
|
||||
command,
|
||||
client_name: ENVIRONMENT_CLIENT_NAME.to_string(),
|
||||
initialize_timeout,
|
||||
resume_session_id: None,
|
||||
})
|
||||
.await
|
||||
.await;
|
||||
}
|
||||
}
|
||||
};
|
||||
Self::connect_websocket(RemoteExecServerConnectArgs {
|
||||
websocket_url,
|
||||
client_name: ENVIRONMENT_CLIENT_NAME.to_string(),
|
||||
connect_timeout,
|
||||
initialize_timeout,
|
||||
resume_session_id: None,
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn open_initial_noise_rendezvous_connection(
|
||||
|
||||
@@ -3,6 +3,8 @@ use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use futures::FutureExt;
|
||||
|
||||
use crate::ExecServerError;
|
||||
use crate::ExecServerRuntimePaths;
|
||||
use crate::ExecutorFileSystem;
|
||||
@@ -26,6 +28,7 @@ use crate::protocol::EnvironmentInfo;
|
||||
use crate::remote::NoiseRendezvousEnvironmentConfig;
|
||||
use crate::remote_file_system::RemoteFileSystem;
|
||||
use crate::remote_process::RemoteProcess;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio_util::task::AbortOnDropHandle;
|
||||
|
||||
pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL";
|
||||
@@ -59,6 +62,10 @@ pub struct EnvironmentManager {
|
||||
local_runtime_paths: Option<ExecServerRuntimePaths>,
|
||||
}
|
||||
|
||||
/// The one-shot capability to complete a pending environment registration.
|
||||
#[must_use = "the pending environment cannot connect until registration is completed"]
|
||||
pub struct PendingEnvironmentRegistration(oneshot::Sender<Result<String, String>>);
|
||||
|
||||
pub const LOCAL_ENVIRONMENT_ID: &str = "local";
|
||||
pub const REMOTE_ENVIRONMENT_ID: &str = "remote";
|
||||
|
||||
@@ -295,22 +302,8 @@ impl EnvironmentManager {
|
||||
exec_server_url: String,
|
||||
connect_timeout: Option<std::time::Duration>,
|
||||
) -> Result<(), ExecServerError> {
|
||||
if environment_id.is_empty() {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"environment id cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
let (exec_server_url, disabled) = normalize_exec_server_url(Some(exec_server_url));
|
||||
if disabled {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"remote environment cannot use disabled exec-server url".to_string(),
|
||||
));
|
||||
}
|
||||
let Some(exec_server_url) = exec_server_url else {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"remote environment requires an exec-server url".to_string(),
|
||||
));
|
||||
};
|
||||
validate_environment_id(&environment_id)?;
|
||||
let exec_server_url = validate_remote_exec_server_url(exec_server_url)?;
|
||||
let environment = Arc::new(Environment::remote_with_transport(
|
||||
ExecServerTransportParams::websocket_url(
|
||||
exec_server_url,
|
||||
@@ -326,6 +319,25 @@ impl EnvironmentManager {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Adds or replaces a remote environment whose stable URL will be supplied later.
|
||||
pub fn register_pending_environment(
|
||||
&self,
|
||||
environment_id: String,
|
||||
) -> Result<PendingEnvironmentRegistration, ExecServerError> {
|
||||
validate_environment_id(&environment_id)?;
|
||||
let (completion, websocket_url) = oneshot::channel();
|
||||
let environment = Arc::new(Environment::remote_with_transport(
|
||||
ExecServerTransportParams::PendingWebSocketUrl(websocket_url.shared()),
|
||||
self.local_runtime_paths.clone(),
|
||||
));
|
||||
environment.start_connecting();
|
||||
self.environments
|
||||
.write()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.insert(environment_id, environment);
|
||||
Ok(PendingEnvironmentRegistration(completion))
|
||||
}
|
||||
|
||||
/// Adds or replaces a named remote environment that connects through an
|
||||
/// authenticated, end-to-end encrypted rendezvous stream.
|
||||
///
|
||||
@@ -336,11 +348,7 @@ impl EnvironmentManager {
|
||||
environment_id: String,
|
||||
provider: Arc<dyn NoiseRendezvousConnectProvider>,
|
||||
) -> Result<(), ExecServerError> {
|
||||
if environment_id.is_empty() {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"environment id cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
validate_environment_id(&environment_id)?;
|
||||
let identity = NoiseChannelIdentity::generate().map_err(|error| {
|
||||
ExecServerError::Protocol(format!(
|
||||
"failed to generate Noise harness identity: {error}"
|
||||
@@ -359,6 +367,46 @@ impl EnvironmentManager {
|
||||
}
|
||||
}
|
||||
|
||||
impl PendingEnvironmentRegistration {
|
||||
/// Completes provisioning with the stable URL or a terminal error message.
|
||||
pub fn complete(self, result: Result<String, String>) -> Result<(), ExecServerError> {
|
||||
let result = match result {
|
||||
Ok(exec_server_url) => match validate_remote_exec_server_url(exec_server_url) {
|
||||
Ok(exec_server_url) => Ok(exec_server_url),
|
||||
Err(error) => {
|
||||
let _ = self.0.send(Err(error.to_string()));
|
||||
return Err(error);
|
||||
}
|
||||
},
|
||||
Err(message) => Err(message),
|
||||
};
|
||||
self.0.send(result).map_err(|_| {
|
||||
ExecServerError::Disconnected("pending environment registration is inactive".into())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_environment_id(environment_id: &str) -> Result<(), ExecServerError> {
|
||||
if environment_id.is_empty() {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"environment id cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_remote_exec_server_url(exec_server_url: String) -> Result<String, ExecServerError> {
|
||||
let (exec_server_url, disabled) = normalize_exec_server_url(Some(exec_server_url));
|
||||
if disabled {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"remote environment cannot use disabled exec-server url".to_string(),
|
||||
));
|
||||
}
|
||||
exec_server_url.ok_or_else(|| {
|
||||
ExecServerError::Protocol("remote environment requires an exec-server url".to_string())
|
||||
})
|
||||
}
|
||||
|
||||
fn noise_environment_config_from_env()
|
||||
-> Result<Option<NoiseRendezvousEnvironmentConfig>, ExecServerError> {
|
||||
noise_environment_config_from_values(
|
||||
@@ -412,7 +460,6 @@ fn optional_environment_value(name: &str) -> Option<String> {
|
||||
/// paths used by filesystem helpers.
|
||||
#[derive(Clone)]
|
||||
pub struct Environment {
|
||||
exec_server_url: Option<String>,
|
||||
remote_client: Option<LazyRemoteExecServerClient>,
|
||||
// Dropping the environment stops unfinished background startup work.
|
||||
startup_task: Arc<Mutex<Option<AbortOnDropHandle<()>>>>,
|
||||
@@ -426,7 +473,6 @@ impl Environment {
|
||||
/// Builds a test-only local environment without configured sandbox helper paths.
|
||||
pub fn default_for_tests() -> Self {
|
||||
Self {
|
||||
exec_server_url: None,
|
||||
remote_client: None,
|
||||
startup_task: Arc::new(Mutex::new(None)),
|
||||
exec_backend: Arc::new(LocalProcess::default()),
|
||||
@@ -439,9 +485,7 @@ impl Environment {
|
||||
|
||||
impl std::fmt::Debug for Environment {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("Environment")
|
||||
.field("exec_server_url", &self.exec_server_url)
|
||||
.finish_non_exhaustive()
|
||||
f.debug_struct("Environment").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -483,7 +527,6 @@ impl Environment {
|
||||
|
||||
pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self {
|
||||
Self {
|
||||
exec_server_url: None,
|
||||
remote_client: None,
|
||||
startup_task: Arc::new(Mutex::new(None)),
|
||||
exec_backend: Arc::new(LocalProcess::with_local_runtime_paths(
|
||||
@@ -514,21 +557,12 @@ impl Environment {
|
||||
remote_transport: ExecServerTransportParams,
|
||||
local_runtime_paths: Option<ExecServerRuntimePaths>,
|
||||
) -> Self {
|
||||
let exec_server_url = match &remote_transport {
|
||||
ExecServerTransportParams::WebSocketUrl {
|
||||
websocket_url: exec_server_url,
|
||||
..
|
||||
} => Some(exec_server_url.clone()),
|
||||
ExecServerTransportParams::NoiseRendezvous { .. } => None,
|
||||
ExecServerTransportParams::StdioCommand { .. } => None,
|
||||
};
|
||||
let client = LazyRemoteExecServerClient::new(remote_transport);
|
||||
let exec_backend: Arc<dyn ExecBackend> = Arc::new(RemoteProcess::new(client.clone()));
|
||||
let filesystem: Arc<dyn ExecutorFileSystem> =
|
||||
Arc::new(RemoteFileSystem::new(client.clone()));
|
||||
|
||||
Self {
|
||||
exec_server_url,
|
||||
remote_client: Some(client.clone()),
|
||||
startup_task: Arc::new(Mutex::new(None)),
|
||||
exec_backend,
|
||||
@@ -542,11 +576,6 @@ impl Environment {
|
||||
self.remote_client.is_some()
|
||||
}
|
||||
|
||||
/// Returns the remote exec-server URL when this environment is remote.
|
||||
pub fn exec_server_url(&self) -> Option<&str> {
|
||||
self.exec_server_url.as_deref()
|
||||
}
|
||||
|
||||
pub fn local_runtime_paths(&self) -> Option<&ExecServerRuntimePaths> {
|
||||
self.local_runtime_paths.as_ref()
|
||||
}
|
||||
@@ -718,7 +747,6 @@ mod tests {
|
||||
let environment = Environment::create(/*exec_server_url*/ None, test_runtime_paths())
|
||||
.expect("create environment");
|
||||
|
||||
assert_eq!(environment.exec_server_url(), None);
|
||||
assert!(!environment.is_remote());
|
||||
assert!(environment.info().await.is_ok());
|
||||
}
|
||||
@@ -758,7 +786,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn environment_manager_reports_remote_url() {
|
||||
async fn environment_manager_creates_remote_environment_for_url() {
|
||||
let manager = EnvironmentManager::create_for_tests(
|
||||
Some("ws://127.0.0.1:8765".to_string()),
|
||||
Some(test_runtime_paths()),
|
||||
@@ -771,7 +799,6 @@ mod tests {
|
||||
Some(REMOTE_ENVIRONMENT_ID)
|
||||
);
|
||||
assert!(environment.is_remote());
|
||||
assert_eq!(environment.exec_server_url(), Some("ws://127.0.0.1:8765"));
|
||||
assert!(Arc::ptr_eq(
|
||||
&environment,
|
||||
&manager
|
||||
@@ -961,7 +988,7 @@ mod tests {
|
||||
|
||||
assert_eq!(environment.local_runtime_paths(), Some(&runtime_paths));
|
||||
let manager = EnvironmentManager::create_for_tests(
|
||||
environment.exec_server_url().map(str::to_owned),
|
||||
/*exec_server_url*/ None,
|
||||
Some(
|
||||
environment
|
||||
.local_runtime_paths()
|
||||
@@ -1030,7 +1057,6 @@ mod tests {
|
||||
.get_environment("executor-a")
|
||||
.expect("first remote environment");
|
||||
assert!(first.is_remote());
|
||||
assert_eq!(first.exec_server_url(), Some("ws://127.0.0.1:8765"));
|
||||
assert_eq!(manager.default_environment_id(), None);
|
||||
|
||||
manager
|
||||
@@ -1044,7 +1070,6 @@ mod tests {
|
||||
.get_environment("executor-a")
|
||||
.expect("second remote environment");
|
||||
assert!(second.is_remote());
|
||||
assert_eq!(second.exec_server_url(), Some("ws://127.0.0.1:9876"));
|
||||
assert!(!Arc::ptr_eq(&first, &second));
|
||||
}
|
||||
|
||||
|
||||
@@ -177,25 +177,17 @@ mod tests {
|
||||
assert!(!environments.contains_key(LOCAL_ENVIRONMENT_ID));
|
||||
let remote_environment = &environments[REMOTE_ENVIRONMENT_ID];
|
||||
assert!(remote_environment.is_remote());
|
||||
assert_eq!(
|
||||
remote_environment.exec_server_url(),
|
||||
Some("ws://127.0.0.1:8765")
|
||||
);
|
||||
assert_eq!(
|
||||
default,
|
||||
EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_provider_normalizes_exec_server_url() {
|
||||
let provider = DefaultEnvironmentProvider::new(Some(" ws://127.0.0.1:8765 ".to_string()));
|
||||
let snapshot = provider.snapshot().await.expect("environments");
|
||||
let environments: HashMap<_, _> = snapshot.environments.into_iter().collect();
|
||||
|
||||
#[test]
|
||||
fn normalizes_exec_server_url() {
|
||||
assert_eq!(
|
||||
environments[REMOTE_ENVIRONMENT_ID].exec_server_url(),
|
||||
Some("ws://127.0.0.1:8765")
|
||||
normalize_exec_server_url(Some(" ws://127.0.0.1:8765 ".to_string())),
|
||||
(Some("ws://127.0.0.1:8765".to_string()), false)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -389,12 +389,8 @@ mod tests {
|
||||
|
||||
assert!(include_local);
|
||||
assert!(!environments.contains_key(LOCAL_ENVIRONMENT_ID));
|
||||
assert_eq!(
|
||||
environments["devbox"].exec_server_url(),
|
||||
Some("ws://127.0.0.1:8765")
|
||||
);
|
||||
assert!(environments["devbox"].is_remote());
|
||||
assert!(environments["ssh-dev"].is_remote());
|
||||
assert_eq!(environments["ssh-dev"].exec_server_url(), None);
|
||||
assert_eq!(
|
||||
default,
|
||||
EnvironmentDefault::EnvironmentId("ssh-dev".to_string())
|
||||
|
||||
@@ -68,6 +68,7 @@ pub use environment::CODEX_EXEC_SERVER_URL_ENV_VAR;
|
||||
pub use environment::Environment;
|
||||
pub use environment::EnvironmentManager;
|
||||
pub use environment::LOCAL_ENVIRONMENT_ID;
|
||||
pub use environment::PendingEnvironmentRegistration;
|
||||
pub use environment::REMOTE_ENVIRONMENT_ID;
|
||||
pub use environment_provider::DefaultEnvironmentProvider;
|
||||
pub use environment_provider::EnvironmentProvider;
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
mod common;
|
||||
|
||||
use codex_exec_server::EnvironmentManager;
|
||||
use common::exec_server::exec_server;
|
||||
use futures::poll;
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn pending_environment_connects_and_reconnects_after_completion() -> anyhow::Result<()> {
|
||||
let mut server = exec_server().await?;
|
||||
let mut proxy = server.disconnectable_websocket_proxy().await?;
|
||||
let manager = EnvironmentManager::without_environments();
|
||||
let registration = manager.register_pending_environment("tools".to_string())?;
|
||||
let environment = manager.get_environment("tools").expect("environment");
|
||||
|
||||
registration.complete(Ok(proxy.websocket_url().to_string()))?;
|
||||
environment.wait_until_ready().await?;
|
||||
proxy.pause_and_disconnect().await?;
|
||||
proxy.resume()?;
|
||||
environment.info().await?;
|
||||
server.shutdown().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failure_and_dropped_registration_are_terminal() -> anyhow::Result<()> {
|
||||
let manager = EnvironmentManager::without_environments();
|
||||
let failed = manager.register_pending_environment("failed".to_string())?;
|
||||
let failed_environment = manager.get_environment("failed").expect("environment");
|
||||
failed.complete(Err("provisioning failed".to_string()))?;
|
||||
let error = failed_environment.wait_until_ready().await.unwrap_err();
|
||||
let message = error.to_string();
|
||||
assert!(message.ends_with("environment unavailable: provisioning failed"));
|
||||
|
||||
let dropped = manager.register_pending_environment("dropped".to_string())?;
|
||||
let dropped_environment = manager.get_environment("dropped").expect("environment");
|
||||
drop(dropped);
|
||||
let error = dropped_environment.wait_until_ready().await.unwrap_err();
|
||||
let message = error.to_string();
|
||||
assert!(message.contains("registration ended before completion"));
|
||||
assert!(manager.get_environment("failed").is_some());
|
||||
assert!(manager.get_environment("dropped").is_some());
|
||||
|
||||
let invalid = manager.register_pending_environment("invalid".to_string())?;
|
||||
let invalid_environment = manager.get_environment("invalid").expect("environment");
|
||||
let error = invalid.complete(Ok(String::new())).unwrap_err();
|
||||
assert!(error.to_string().contains("requires an exec-server url"));
|
||||
let error = invalid_environment.wait_until_ready().await.unwrap_err();
|
||||
assert!(error.to_string().contains("requires an exec-server url"));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn late_completion_is_isolated_from_replacement() -> anyhow::Result<()> {
|
||||
let mut server = exec_server().await?;
|
||||
let manager = EnvironmentManager::without_environments();
|
||||
let old_registration = manager.register_pending_environment("tools".to_string())?;
|
||||
let old_environment = manager.get_environment("tools").expect("old environment");
|
||||
let current_registration = manager.register_pending_environment("tools".to_string())?;
|
||||
let current = manager.get_environment("tools").expect("current");
|
||||
|
||||
old_registration.complete(Ok(server.websocket_url().to_string()))?;
|
||||
old_environment.wait_until_ready().await?;
|
||||
let mut current_readiness = Box::pin(current.wait_until_ready());
|
||||
assert!(poll!(&mut current_readiness).is_pending());
|
||||
|
||||
current_registration.complete(Ok(server.websocket_url().to_string()))?;
|
||||
current_readiness.await?;
|
||||
server.shutdown().await?;
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user