openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-06 15:29:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

Defer Noise environment connections until registration (#33166)

Browse Source 
## What changed

- Replace pending WebSocket URL registration with deferred Noise environment
  registration that gates connection attempts on an explicit readiness signal.
- Reuse the Noise rendezvous transport after readiness so reconnects request a
  fresh connection bundle.
- Preserve terminal errors for failed or dropped registrations and keep late
  completion isolated from replacement environments.

## Testing

- Add coverage for readiness gating, registration failure and replacement,
  eager Noise connections, and reconnection through a fresh rendezvous bundle.

GitOrigin-RevId: 83e23fa03a02e3b2bdf1a83fe26d7ac461f55cf3
This commit is contained in:
TAFOYA-OAI
2026-07-14 21:02:20 +00:00
committed by copyberry
parent 5d325ba223
commit 32cd5d4eab
9 changed files with 369 additions and 216 deletions
Download Patch File Download Diff File Expand all files Collapse all files

138
codex-rs/core/tests/suite/remote_env.rs
View File

@@ -7,8 +7,12 @@ use codex_core::compact::SUMMARIZATION_PROMPT;
use codex_core::config::Constrained;
use codex_exec_server::CopyOptions;
use codex_exec_server::CreateDirectoryOptions;
use codex_exec_server::ExecServerError;
use codex_exec_server::FileSystemSandboxContext;
use codex_exec_server::LOCAL_ENVIRONMENT_ID;
use codex_exec_server::NoiseChannelPublicKey;
use codex_exec_server::NoiseRendezvousConnectBundle;
use codex_exec_server::NoiseRendezvousConnectProvider;
use codex_exec_server::REMOTE_ENVIRONMENT_ID;
use codex_exec_server::RemoveOptions;
use codex_features::Feature;
@@ -69,6 +73,7 @@ use core_test_support::wait_for_event;
use core_test_support::wait_for_event_match;
use futures::SinkExt;
use futures::StreamExt;
use futures::future::BoxFuture;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
@@ -76,6 +81,9 @@ use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use std::process::Command;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::SystemTime;
use std::time::UNIX_EPOCH;
@@ -483,6 +491,25 @@ fn tool_names(body: &Value) -> Vec<String> {
.collect()
}
#[derive(Default)]
struct FailingNoiseConnectProvider {
calls: AtomicUsize,
}
impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider {
fn connect_bundle(
&self,
_: NoiseChannelPublicKey,
) -> BoxFuture<'_, std::result::Result<NoiseRendezvousConnectBundle, ExecServerError>> {
self.calls.fetch_add(1, Ordering::Relaxed);
Box::pin(async {
Err(ExecServerError::Protocol(
"test Noise connection failed".to_string(),
))
})
}
}
async fn wait_for_response_request_count(response_mock: &ResponseMock, expected_count: usize) {
timeout(Duration::from_secs(5), async {
while response_mock.requests().len() < expected_count {
@@ -494,8 +521,7 @@ async fn wait_for_response_request_count(response_mock: &ResponseMock, expected_
}
#[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?;
async fn deferred_executor_starts_noise_connection_after_registration() -> Result<()> {
let server = start_mock_server().await;
let wait_call_id = "wait-for-startup";
let response_mock = mount_sse_sequence(
@@ -515,47 +541,24 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<(
]),
sse(vec![
ev_response_created("resp-2"),
ev_function_call(
"request-permissions",
"request_permissions",
&json!({
"reason": "Verify that the ready environment is used.",
"permissions": {
"network": { "enabled": true }
}
})
.to_string(),
),
ev_assistant_message("msg-2", "done"),
ev_completed("resp-2"),
]),
sse(vec![
ev_response_created("resp-3"),
ev_assistant_message("msg-3", "done"),
ev_completed("resp-3"),
]),
],
)
.await;
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())?;
let provider = Arc::new(FailingNoiseConnectProvider::default());
let registration = environment_manager
.register_deferred_noise_environment(REMOTE_ENVIRONMENT_ID.to_string(), provider.clone())?;
test.codex
.submit(Op::UserInput {
@@ -581,83 +584,26 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<(
.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!(
event,
EventMsg::RequestPermissions(_) | EventMsg::TurnComplete(_)
)
})
.await;
let EventMsg::RequestPermissions(permission_request) = event else {
panic!("ready environment should be available to request_permissions: {event:?}");
};
assert_eq!(
permission_request.environment_id.as_deref(),
Some(REMOTE_ENVIRONMENT_ID)
);
test.codex
.submit(Op::RequestPermissionsResponse {
id: permission_request.call_id,
response: RequestPermissionsResponse {
permissions: RequestPermissionProfile::default(),
scope: PermissionGrantScope::Turn,
strict_auto_review: false,
},
})
.await?;
assert_eq!(provider.calls.load(Ordering::Relaxed), 0);
registration.complete(Ok(()))?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
assert_eq!(provider.calls.load(Ordering::Relaxed), 1);
let requests = response_mock.requests();
assert_eq!(requests.len(), 3);
assert_eq!(requests.len(), 2);
let starting_tools = tool_names(&requests[0].body_json());
let ready_tools = tool_names(&requests[1].body_json());
assert!(starting_tools.contains(&"wait_for_environment".to_string()));
assert!(!starting_tools.contains(&"exec_command".to_string()));
assert!(ready_tools.contains(&"exec_command".to_string()));
assert!(ready_tools.contains(&"wait_for_environment".to_string()));
let (wait_output, _) = requests[1]
.function_call_output_content_and_success(wait_call_id)
.context("wait_for_environment output should be present")?;
assert_eq!(
serde_json::from_str::<Value>(&wait_output.context("wait output should contain text")?)?,
json!({
"environment_id": REMOTE_ENVIRONMENT_ID,
"status": "ready",
})
);
assert!(
requests[0]
.message_input_texts("user")
.iter()
.any(|text| text.contains("<status>starting</status>"))
);
let ready_user_context = requests[1].message_input_texts("user");
assert_eq!(
ready_user_context
.iter()
.filter(|text| text.contains("<shell>zsh</shell>"))
.count(),
1
);
let final_user_context = requests[2].message_input_texts("user");
assert_eq!(
final_user_context
.iter()
.filter(|text| text.contains("<status>starting</status>"))
.count(),
1
);
assert_eq!(
final_user_context
.iter()
.filter(|text| text.contains("<shell>zsh</shell>"))
.count(),
1
wait_output
.context("wait output should contain text")?
.contains("failed to start")
);
Ok(())
@@ -792,8 +738,10 @@ async fn deferred_executor_wait_reports_startup_failure() -> Result<()> {
.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())?;
let registration = environment_manager.register_deferred_noise_environment(
REMOTE_ENVIRONMENT_ID.to_string(),
Arc::new(FailingNoiseConnectProvider::default()),
)?;
test.codex
.submit(Op::UserInput {

4
codex-rs/exec-server/src/client.rs
View File

@@ -453,8 +453,8 @@ impl LazyRemoteExecServerClient {
fn can_reconnect(&self) -> bool {
matches!(
self.transport_params,
ExecServerTransportParams::WebSocketUrl { .. }
| ExecServerTransportParams::PendingWebSocketUrl(..)
ExecServerTransportParams::Deferred(_)
| ExecServerTransportParams::WebSocketUrl { .. }
| ExecServerTransportParams::NoiseRendezvous { .. }
)
}

17
codex-rs/exec-server/src/client_api.rs
View File

@@ -90,17 +90,23 @@ pub(crate) struct StdioExecServerCommand {
pub cwd: Option<PathBuf>,
}
pub(crate) type PendingExecServerUrl = Shared<oneshot::Receiver<Result<String, String>>>;
pub(crate) type DeferredEnvironmentReadiness = Shared<oneshot::Receiver<Result<(), String>>>;
#[derive(Clone)]
pub(crate) struct Deferred<T> {
pub readiness: DeferredEnvironmentReadiness,
pub transport: T,
}
/// Parameters used to connect to a remote exec-server environment.
#[derive(Clone)]
pub(crate) enum ExecServerTransportParams {
Deferred(Box<Deferred<ExecServerTransportParams>>),
WebSocketUrl {
websocket_url: String,
connect_timeout: Duration,
initialize_timeout: Duration,
},
PendingWebSocketUrl(PendingExecServerUrl),
NoiseRendezvous {
provider: Arc<dyn NoiseRendezvousConnectProvider>,
identity: NoiseChannelIdentity,
@@ -115,6 +121,10 @@ pub(crate) enum ExecServerTransportParams {
impl std::fmt::Debug for ExecServerTransportParams {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Deferred(deferred) => f
.debug_struct("Deferred")
.field("transport", &deferred.transport)
.finish_non_exhaustive(),
Self::WebSocketUrl {
websocket_url,
connect_timeout,
@@ -125,9 +135,6 @@ 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()
}

38
codex-rs/exec-server/src/client_transport.rs
View File

@@ -94,27 +94,35 @@ impl ExecServerClient {
pub(crate) async fn connect_for_transport(
transport_params: ExecServerTransportParams,
) -> Result<Self, ExecServerError> {
let (transport_params, deferred_readiness) = match transport_params {
ExecServerTransportParams::Deferred(deferred) => {
(deferred.transport, Some(deferred.readiness))
}
transport_params => (transport_params, None),
};
if let Some(readiness) = deferred_readiness {
readiness
.await
.unwrap_or_else(|_| {
Err("environment registration ended before completion".to_string())
})
.map_err(|message| {
ExecServerError::Disconnected(format!("environment unavailable: {message}"))
})?;
}
let (websocket_url, connect_timeout, initialize_timeout) = match transport_params {
ExecServerTransportParams::Deferred(_) => {
return Err(ExecServerError::Protocol(
"nested deferred exec-server transports are unsupported".to_string(),
));
}
ExecServerTransportParams::WebSocketUrl {
websocket_url,
connect_timeout,
initialize_timeout,
} => (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,
DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
)
}
ExecServerTransportParams::NoiseRendezvous { provider, identity } => {
let reconnect_strategy = ExecServerReconnectStrategy::NoiseRendezvous {
provider: Arc::clone(&provider),

55
codex-rs/exec-server/src/environment.rs
View File

@@ -62,9 +62,9 @@ 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>>);
/// The one-shot capability to complete a deferred environment registration.
#[must_use = "the deferred environment cannot connect until registration is completed"]
pub struct DeferredEnvironmentRegistration(oneshot::Sender<Result<(), String>>);
pub const LOCAL_ENVIRONMENT_ID: &str = "local";
pub const REMOTE_ENVIRONMENT_ID: &str = "remote";
@@ -352,15 +352,20 @@ impl EnvironmentManager {
Ok(())
}
/// Adds or replaces a remote environment whose stable URL will be supplied later.
pub fn register_pending_environment(
/// Adds or replaces a Noise rendezvous environment that will become ready later.
pub fn register_deferred_noise_environment(
&self,
environment_id: String,
) -> Result<PendingEnvironmentRegistration, ExecServerError> {
provider: Arc<dyn NoiseRendezvousConnectProvider>,
) -> Result<DeferredEnvironmentRegistration, ExecServerError> {
validate_environment_id(&environment_id)?;
let (completion, websocket_url) = oneshot::channel();
let identity = noise_channel_identity()?;
let (completion, readiness) = oneshot::channel();
let environment = Arc::new(Environment::remote_with_transport(
ExecServerTransportParams::PendingWebSocketUrl(websocket_url.shared()),
ExecServerTransportParams::Deferred(Box::new(crate::client_api::Deferred {
readiness: readiness.shared(),
transport: ExecServerTransportParams::NoiseRendezvous { provider, identity },
})),
self.local_runtime_paths.clone(),
));
environment.start_connecting();
@@ -368,7 +373,7 @@ impl EnvironmentManager {
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(environment_id, environment);
Ok(PendingEnvironmentRegistration(completion))
Ok(DeferredEnvironmentRegistration(completion))
}
/// Adds or replaces a named remote environment that connects through an
@@ -382,11 +387,7 @@ impl EnvironmentManager {
provider: Arc<dyn NoiseRendezvousConnectProvider>,
) -> Result<(), ExecServerError> {
validate_environment_id(&environment_id)?;
let identity = NoiseChannelIdentity::generate().map_err(|error| {
ExecServerError::Protocol(format!(
"failed to generate Noise harness identity: {error}"
))
})?;
let identity = noise_channel_identity()?;
let environment = Arc::new(Environment::remote_with_transport(
ExecServerTransportParams::NoiseRendezvous { provider, identity },
self.local_runtime_paths.clone(),
@@ -400,25 +401,23 @@ 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),
};
impl DeferredEnvironmentRegistration {
/// Completes provisioning with readiness or a terminal error message.
pub fn complete(self, result: Result<(), String>) -> Result<(), ExecServerError> {
self.0.send(result).map_err(|_| {
ExecServerError::Disconnected("pending environment registration is inactive".into())
ExecServerError::Disconnected("deferred environment registration is inactive".into())
})
}
}
fn noise_channel_identity() -> Result<NoiseChannelIdentity, ExecServerError> {
NoiseChannelIdentity::generate().map_err(|error| {
ExecServerError::Protocol(format!(
"failed to generate Noise harness identity: {error}"
))
})
}
fn validate_environment_id(environment_id: &str) -> Result<(), ExecServerError> {
if environment_id.is_empty() {
return Err(ExecServerError::Protocol(

2
codex-rs/exec-server/src/lib.rs
View File

@@ -65,11 +65,11 @@ pub use environment::CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID_ENV_VAR;
pub use environment::CODEX_EXEC_SERVER_NOISE_ENVIRONMENT_ID_ENV_VAR;
pub use environment::CODEX_EXEC_SERVER_NOISE_REGISTRY_URL_ENV_VAR;
pub use environment::CODEX_EXEC_SERVER_URL_ENV_VAR;
pub use environment::DeferredEnvironmentRegistration;
pub use environment::Environment;
pub use environment::EnvironmentManager;
pub use environment::EnvironmentObservedStatus;
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;

135
codex-rs/exec-server/tests/deferred_environment.rs Normal file
View File

@@ -0,0 +1,135 @@
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::ExecServerError;
use codex_exec_server::NoiseChannelPublicKey;
use codex_exec_server::NoiseRendezvousConnectBundle;
use codex_exec_server::NoiseRendezvousConnectProvider;
use futures::FutureExt;
use futures::future::BoxFuture;
use futures::poll;
use pretty_assertions::assert_eq;
#[derive(Default)]
struct FailingNoiseConnectProvider {
calls: AtomicUsize,
}
impl FailingNoiseConnectProvider {
fn calls(&self) -> usize {
self.calls.load(Ordering::Relaxed)
}
}
impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider {
fn connect_bundle(
&self,
_: NoiseChannelPublicKey,
) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
self.calls.fetch_add(1, Ordering::Relaxed);
async {
Err(ExecServerError::Protocol(
"test Noise provider called".to_string(),
))
}
.boxed()
}
}
#[tokio::test]
async fn deferred_environment_waits_before_connecting() -> anyhow::Result<()> {
let manager = EnvironmentManager::without_environments();
let provider = Arc::new(FailingNoiseConnectProvider::default());
let registration =
manager.register_deferred_noise_environment("tools".to_string(), provider.clone())?;
let environment = manager.get_environment("tools").expect("environment");
let mut readiness = Box::pin(environment.wait_until_ready());
assert!(poll!(&mut readiness).is_pending());
assert_eq!(provider.calls(), 0);
registration.complete(Ok(()))?;
let error = readiness.await.unwrap_err();
assert!(error.to_string().contains("test Noise provider called"));
assert_eq!(provider.calls(), 1);
Ok(())
}
#[tokio::test]
async fn failure_and_dropped_registration_are_terminal() -> anyhow::Result<()> {
let manager = EnvironmentManager::without_environments();
let failed_provider = Arc::new(FailingNoiseConnectProvider::default());
let failed = manager
.register_deferred_noise_environment("failed".to_string(), failed_provider.clone())?;
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();
assert!(
error
.to_string()
.ends_with("environment unavailable: provisioning failed")
);
assert_eq!(failed_provider.calls(), 0);
let dropped_provider = Arc::new(FailingNoiseConnectProvider::default());
let dropped = manager
.register_deferred_noise_environment("dropped".to_string(), dropped_provider.clone())?;
let dropped_environment = manager.get_environment("dropped").expect("environment");
drop(dropped);
let error = dropped_environment.wait_until_ready().await.unwrap_err();
assert!(
error
.to_string()
.contains("registration ended before completion")
);
assert_eq!(dropped_provider.calls(), 0);
assert!(manager.get_environment("failed").is_some());
assert!(manager.get_environment("dropped").is_some());
Ok(())
}
#[tokio::test]
async fn late_completion_is_isolated_from_replacement() -> anyhow::Result<()> {
let manager = EnvironmentManager::without_environments();
let old_provider = Arc::new(FailingNoiseConnectProvider::default());
let old_registration =
manager.register_deferred_noise_environment("tools".to_string(), old_provider.clone())?;
let old_environment = manager.get_environment("tools").expect("old environment");
let current_provider = Arc::new(FailingNoiseConnectProvider::default());
let current_registration = manager
.register_deferred_noise_environment("tools".to_string(), current_provider.clone())?;
let current = manager.get_environment("tools").expect("current");
old_registration.complete(Ok(()))?;
let old_error = old_environment.wait_until_ready().await.unwrap_err();
assert!(old_error.to_string().contains("test Noise provider called"));
assert_eq!(old_provider.calls(), 1);
let mut current_readiness = Box::pin(current.wait_until_ready());
assert!(poll!(&mut current_readiness).is_pending());
assert_eq!(current_provider.calls(), 0);
current_registration.complete(Ok(()))?;
let current_error = current_readiness.await.unwrap_err();
assert!(
current_error
.to_string()
.contains("test Noise provider called")
);
assert_eq!(current_provider.calls(), 1);
Ok(())
}
#[tokio::test]
async fn eager_noise_environment_connects_without_registration() -> anyhow::Result<()> {
let manager = EnvironmentManager::without_environments();
let provider = Arc::new(FailingNoiseConnectProvider::default());
manager.upsert_noise_environment("tools".to_string(), provider.clone())?;
let environment = manager.get_environment("tools").expect("environment");
let error = environment.wait_until_ready().await.unwrap_err();
assert!(error.to_string().contains("test Noise provider called"));
assert_eq!(provider.calls(), 1);
Ok(())
}

70
codex-rs/exec-server/tests/pending_environment.rs
View File

@@ -1,70 +0,0 @@
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(())
}

126
codex-rs/exec-server/tests/relay.rs
View File

@@ -6,6 +6,8 @@ mod relay_proto;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;
use anyhow::Context;
@@ -13,20 +15,24 @@ use anyhow::Result;
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD;
use codex_api::AuthProvider;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::ExecParams;
use codex_exec_server::ExecResponse;
use codex_exec_server::ExecServerClient;
use codex_exec_server::ExecServerError;
use codex_exec_server::ExecServerRuntimePaths;
use codex_exec_server::FsReadFileParams;
use codex_exec_server::NoiseChannelIdentity;
use codex_exec_server::NoiseChannelPublicKey;
use codex_exec_server::NoiseRendezvousConnectArgs;
use codex_exec_server::NoiseRendezvousConnectBundle;
use codex_exec_server::NoiseRendezvousConnectProvider;
use codex_exec_server::ProcessId;
use codex_exec_server::RemoteEnvironmentConfig;
use codex_utils_path_uri::PathUri;
use futures::SinkExt;
use futures::StreamExt;
use futures::future::BoxFuture;
use http::HeaderMap;
use http::HeaderValue;
use pretty_assertions::assert_eq;
@@ -69,6 +75,126 @@ fn static_registry_auth_provider() -> codex_api::SharedAuthProvider {
Arc::new(StaticRegistryAuthProvider)
}
struct FreshBundleNoiseConnectProvider {
websocket_url: String,
executor_public_key: NoiseChannelPublicKey,
calls: AtomicUsize,
}
impl FreshBundleNoiseConnectProvider {
fn calls(&self) -> usize {
self.calls.load(Ordering::Relaxed)
}
}
impl NoiseRendezvousConnectProvider for FreshBundleNoiseConnectProvider {
fn connect_bundle(
&self,
_: NoiseChannelPublicKey,
) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
let call = self.calls.fetch_add(1, Ordering::Relaxed) + 1;
let bundle = NoiseRendezvousConnectBundle {
websocket_url: self.websocket_url.clone(),
environment_id: ENVIRONMENT_ID.to_string(),
executor_registration_id: EXECUTOR_REGISTRATION_ID.to_string(),
executor_public_key: self.executor_public_key.clone(),
harness_key_authorization: format!("{HARNESS_KEY_AUTHORIZATION}-{call}"),
};
Box::pin(async move { Ok(bundle) })
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let rendezvous_url = format!("ws://{}", listener.local_addr()?);
let registry = MockServer::start().await;
Mock::given(method("POST"))
.and(path(format!(
"/cloud/environment/{ENVIRONMENT_ID}/register"
)))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"environment_id": ENVIRONMENT_ID,
"url": format!("{rendezvous_url}/relay?role=environment"),
"security_profile": "noise_hybrid_ik_v1",
"executor_registration_id": EXECUTOR_REGISTRATION_ID,
})))
.expect(1)
.mount(&registry)
.await;
Mock::given(method("POST"))
.and(path(format!(
"/cloud/environment/{ENVIRONMENT_ID}/validate"
)))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"valid": true,
})))
.expect(2)
.mount(&registry)
.await;
let (codex_exe, codex_linux_sandbox_exe) = common::current_test_binary_helper_paths()?;
let runtime_paths = ExecServerRuntimePaths::new(codex_exe, codex_linux_sandbox_exe)?;
let config = RemoteEnvironmentConfig::new(
registry.uri(),
ENVIRONMENT_ID.to_string(),
static_registry_auth_provider(),
)?;
let remote_environment = tokio::spawn(codex_exec_server::run_remote_environment(
config,
runtime_paths,
));
let environment_websocket = accept_websocket(&listener, "environment").await?;
let provider = Arc::new(FreshBundleNoiseConnectProvider {
websocket_url: format!("{rendezvous_url}/relay?role=harness"),
executor_public_key: registered_executor_public_key(&registry).await?,
calls: AtomicUsize::new(0),
});
let manager = EnvironmentManager::without_environments();
let registration = manager
.register_deferred_noise_environment(ENVIRONMENT_ID.to_string(), provider.clone())?;
let environment = manager
.get_environment(ENVIRONMENT_ID)
.context("deferred Noise environment")?;
assert_eq!(provider.calls(), 0);
registration.complete(Ok(()))?;
let harness_websocket = accept_websocket(&listener, "harness").await?;
let first_relay = tokio::spawn(proxy_relay_frames(
environment_websocket,
harness_websocket,
Arc::new(Mutex::new(Vec::new())),
));
let initial_info = timeout(TEST_TIMEOUT, environment.info())
.await
.context("deferred Noise environment should become ready")??;
assert_eq!(provider.calls(), 1);
first_relay.abort();
let _ = first_relay.await;
let first_reconnected_websocket = accept_websocket(&listener, "reconnected peer").await?;
let second_reconnected_websocket = accept_websocket(&listener, "reconnected peer").await?;
let second_relay = tokio::spawn(proxy_relay_frames(
first_reconnected_websocket,
second_reconnected_websocket,
Arc::new(Mutex::new(Vec::new())),
));
let recovered_info = timeout(TEST_TIMEOUT, environment.info())
.await
.context("deferred Noise environment should reconnect")??;
assert_eq!(recovered_info, initial_info);
assert_eq!(provider.calls(), 2);
registry.verify().await;
second_relay.abort();
remote_environment.abort();
let _ = second_relay.await;
let _ = remote_environment.await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;