codex/codex-rs/exec-server/tests/deferred_environment.rs

use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;

use codex_exec_server::EnvironmentReadyInfo;
use codex_exec_server::ExecServerError;
use codex_exec_server::NoiseChannelPublicKey;
use codex_exec_server::NoiseRendezvousConnectBundle;
use codex_exec_server::NoiseRendezvousConnectProvider;
use codex_exec_server_test_support::environment_manager_without_environments;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_utils_path_uri::PathUri;
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()
    }
}

fn ready_info(root_id: &str, environment_id: &str) -> anyhow::Result<EnvironmentReadyInfo> {
    Ok(EnvironmentReadyInfo {
        selected_capability_roots: vec![SelectedCapabilityRoot {
            id: root_id.to_string(),
            location: CapabilityRootLocation::Environment {
                environment_id: environment_id.to_string(),
                path: PathUri::parse("file:///plugins/root")?,
            },
        }],
    })
}

#[tokio::test]
async fn deferred_environment_waits_before_connecting() -> anyhow::Result<()> {
    let manager = environment_manager_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 connection_state = environment
        .subscribe_connection_state()
        .expect("remote environment connection state");
    let mut readiness = Box::pin(environment.wait_until_ready());

    assert!(poll!(&mut readiness).is_pending());
    assert_eq!(provider.calls(), 0);
    assert!(environment.selected_capability_roots().is_empty());

    let ready_info = ready_info("selected-root", "tools")?;
    registration.complete(Ok(ready_info.clone()))?;
    assert_eq!(
        environment.selected_capability_roots(),
        ready_info.selected_capability_roots
    );
    let error = readiness.await.unwrap_err();
    assert!(error.to_string().contains("test Noise provider called"));
    assert_eq!(provider.calls(), 1);
    assert!(!connection_state.has_changed()?);
    Ok(())
}

#[tokio::test]
async fn existing_environment_publishes_readiness_without_replacement() -> anyhow::Result<()> {
    let manager = environment_manager_without_environments();
    let existing_provider = Arc::new(FailingNoiseConnectProvider::default());
    manager.upsert_noise_environment("tools".to_string(), existing_provider)?;
    let existing_environment = manager
        .get_environment("tools")
        .expect("existing environment");

    let ready_info = ready_info("selected-root", "tools")?;
    manager.publish_ready_info("tools", ready_info.clone())?;

    let current_environment = manager
        .get_environment("tools")
        .expect("current environment");
    assert!(Arc::ptr_eq(&existing_environment, &current_environment));
    assert_eq!(
        existing_environment.selected_capability_roots(),
        ready_info.selected_capability_roots
    );
    Ok(())
}

#[tokio::test]
async fn readiness_updates_the_current_environment_after_replacement() -> anyhow::Result<()> {
    let manager = environment_manager_without_environments();
    manager.upsert_noise_environment(
        "tools".to_string(),
        Arc::new(FailingNoiseConnectProvider::default()),
    )?;
    let captured_environment = manager
        .get_environment("tools")
        .expect("captured environment");
    manager.upsert_noise_environment(
        "tools".to_string(),
        Arc::new(FailingNoiseConnectProvider::default()),
    )?;
    let current_environment = manager
        .get_environment("tools")
        .expect("replacement environment");
    let ready_info = ready_info("selected-root", "tools")?;

    manager.publish_ready_info("tools", ready_info.clone())?;

    assert!(!Arc::ptr_eq(&captured_environment, &current_environment));
    assert!(captured_environment.selected_capability_roots().is_empty());
    assert_eq!(
        current_environment.selected_capability_roots(),
        ready_info.selected_capability_roots
    );
    Ok(())
}

#[tokio::test]
async fn publishing_readiness_requires_existing_environment() -> anyhow::Result<()> {
    let manager = environment_manager_without_environments();

    let error = manager
        .publish_ready_info("tools", ready_info("selected-root", "tools")?)
        .unwrap_err();

    assert!(matches!(error, ExecServerError::Protocol(_)));
    assert!(manager.get_environment("tools").is_none());
    Ok(())
}

#[tokio::test]
async fn existing_environment_accepts_matching_readiness() -> anyhow::Result<()> {
    let manager = environment_manager_without_environments();
    manager.upsert_noise_environment(
        "tools".to_string(),
        Arc::new(FailingNoiseConnectProvider::default()),
    )?;
    let environment = manager.get_environment("tools").expect("environment");
    let ready_info = ready_info("selected-root", "tools")?;

    for _ in 0..2 {
        manager.publish_ready_info("tools", ready_info.clone())?;
    }
    assert_eq!(
        environment.selected_capability_roots(),
        ready_info.selected_capability_roots
    );
    Ok(())
}

#[tokio::test]
async fn existing_environment_overwrites_published_readiness() -> anyhow::Result<()> {
    let manager = environment_manager_without_environments();
    manager.upsert_noise_environment(
        "tools".to_string(),
        Arc::new(FailingNoiseConnectProvider::default()),
    )?;
    let environment = manager.get_environment("tools").expect("environment");
    let selected_ready_info = ready_info("selected-root", "tools")?;
    manager.publish_ready_info("tools", selected_ready_info)?;

    let updated_ready_info = ready_info("different-root", "tools")?;
    manager.publish_ready_info("tools", updated_ready_info.clone())?;
    assert_eq!(
        environment.selected_capability_roots(),
        updated_ready_info.selected_capability_roots
    );
    assert!(Arc::ptr_eq(
        &environment,
        &manager.get_environment("tools").expect("environment")
    ));

    manager.publish_ready_info("tools", EnvironmentReadyInfo::default())?;
    assert!(environment.selected_capability_roots().is_empty());
    Ok(())
}

#[tokio::test]
async fn publishing_readiness_before_deferred_completion_preserves_the_gate() -> anyhow::Result<()>
{
    let manager = environment_manager_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());

    manager.publish_ready_info("tools", ready_info("published-root", "tools")?)?;
    assert!(poll!(&mut readiness).is_pending());
    assert_eq!(provider.calls(), 0);

    let completed_ready_info = ready_info("completed-root", "tools")?;
    registration.complete(Ok(completed_ready_info.clone()))?;
    assert_eq!(
        environment.selected_capability_roots(),
        completed_ready_info.selected_capability_roots
    );
    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 existing_environment_rejects_invalid_readiness() -> anyhow::Result<()> {
    let manager = environment_manager_without_environments();
    manager.upsert_noise_environment(
        "tools".to_string(),
        Arc::new(FailingNoiseConnectProvider::default()),
    )?;
    let existing_environment = manager
        .get_environment("tools")
        .expect("existing environment");
    let error = manager
        .publish_ready_info("tools", ready_info("selected-root", "other")?)
        .unwrap_err();

    assert!(matches!(error, ExecServerError::Protocol(_)));
    assert!(existing_environment.selected_capability_roots().is_empty());
    let current_environment = manager
        .get_environment("tools")
        .expect("current environment");
    assert!(Arc::ptr_eq(&existing_environment, &current_environment));
    Ok(())
}

#[tokio::test]
async fn failure_and_dropped_registration_are_terminal() -> anyhow::Result<()> {
    let manager = environment_manager_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 invalid_ready_info_is_terminal() -> anyhow::Result<()> {
    let manager = environment_manager_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 error = registration
        .complete(Ok(ready_info("selected-root", "other")?))
        .unwrap_err();
    assert!(matches!(error, ExecServerError::Protocol(_)));
    let readiness_error = environment.wait_until_ready().await.unwrap_err();
    assert!(
        readiness_error
            .to_string()
            .contains("belong to environment")
    );
    assert!(environment.selected_capability_roots().is_empty());
    assert_eq!(provider.calls(), 0);
    Ok(())
}

#[tokio::test]
async fn late_completion_is_isolated_from_replacement() -> anyhow::Result<()> {
    let manager = environment_manager_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");

    let old_ready_info = ready_info("old-root", "tools")?;
    old_registration.complete(Ok(old_ready_info.clone()))?;
    assert_eq!(
        old_environment.selected_capability_roots(),
        old_ready_info.selected_capability_roots
    );
    assert!(current.selected_capability_roots().is_empty());
    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);

    let current_ready_info = ready_info("current-root", "tools")?;
    current_registration.complete(Ok(current_ready_info.clone()))?;
    assert_eq!(
        current.selected_capability_roots(),
        current_ready_info.selected_capability_roots
    );
    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 = environment_manager_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(())
}