Files
codex/codex-rs/exec-server/tests/pending_environment.rs
pakrym-oai 4472698728 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
2026-07-10 18:08:35 +00:00

71 lines
3.3 KiB
Rust

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(())
}