github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-27 16:02:59 +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:07:27 +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
codex-rs/core/tests/suite/remote_env.rs
+43 -95
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 {