codex-rs/exec-server/src/client_transport_tests.rs

use std::collections::VecDeque;
use std::future::Future;
use std::sync::Arc;
use std::sync::Mutex;

use anyhow::Result;
use codex_exec_server_protocol::JSONRPCMessage;
use futures::FutureExt;
use futures::SinkExt;
use futures::StreamExt;
use futures::future::BoxFuture;
use pretty_assertions::assert_eq;
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::io::duplex;
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::tungstenite::Message;

use super::ExecServerClient;
use super::ExecServerReconnectStrategy;
use super::INITIAL_REGISTRY_MAX_RETRIES;
use super::INITIAL_REGISTRY_OPERATION_TIMEOUT;
use super::INITIAL_REGISTRY_REQUEST_TIMEOUT;
use crate::ExecServerError;
use crate::NoiseChannelIdentity;
use crate::NoiseChannelPublicKey;
use crate::NoiseRendezvousConnectArgs;
use crate::NoiseRendezvousConnectBundle;
use crate::NoiseRendezvousConnectProvider;
use crate::client::NoiseInitializeContext;
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::connection::JsonRpcConnection;
use crate::noise_channel::PendingResponderHandshake;
use crate::noise_channel::noise_channel_prologue;
use crate::protocol::INITIALIZE_METHOD;
use crate::relay::RelayFrameBodyKind;
use crate::relay::decode_relay_message_frame;
use crate::relay::encode_relay_message_frame;
use crate::relay_proto::RelayMessageFrame;

#[derive(Default)]
struct SequenceNoiseConnectProvider {
    bundles:
        Mutex<VecDeque<BoxFuture<'static, Result<NoiseRendezvousConnectBundle, ExecServerError>>>>,
    returned_urls: Mutex<Vec<String>>,
    requested_keys: Mutex<Vec<NoiseChannelPublicKey>>,
}

impl SequenceNoiseConnectProvider {
    fn push_response(
        &self,
        response: impl Future<Output = Result<NoiseRendezvousConnectBundle, ExecServerError>>
        + Send
        + 'static,
    ) {
        self.bundles.lock().unwrap().push_back(response.boxed());
    }

    fn push_error(&self, error: ExecServerError) {
        self.push_response(futures::future::ready(Err(error)));
    }

    fn push_pending(&self) {
        self.push_response(futures::future::pending());
    }

    fn requested_keys(&self) -> Vec<NoiseChannelPublicKey> {
        self.requested_keys.lock().unwrap().clone()
    }

    fn assert_requested_identity(&self, identity: &NoiseChannelIdentity, requests: usize) {
        assert_eq!(self.requested_keys(), vec![identity.public_key(); requests]);
    }

    fn returned_urls(&self) -> Vec<String> {
        self.returned_urls
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .clone()
    }

    async fn connect(
        self: &Arc<Self>,
        identity: &NoiseChannelIdentity,
    ) -> Result<
        (
            super::JsonRpcConnection,
            super::ExecServerClientConnectOptions,
        ),
        ExecServerError,
    > {
        let provider: Arc<dyn NoiseRendezvousConnectProvider> = self.clone();
        ExecServerClient::open_initial_noise_rendezvous_connection(
            &provider,
            identity,
            codex_http_client::HttpClientFactory::new(
                codex_http_client::OutboundProxyPolicy::ReqwestDefault,
            ),
        )
        .await
        .map(|(ready, _)| (ready.connection, ready.options))
    }
}

impl NoiseRendezvousConnectProvider for SequenceNoiseConnectProvider {
    fn connect_bundle(
        &self,
        harness_public_key: NoiseChannelPublicKey,
    ) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
        self.requested_keys.lock().unwrap().push(harness_public_key);
        let response = self
            .bundles
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .pop_front()
            .expect("test Noise provider exhausted");
        Box::pin(async move {
            let result = response.await;
            if let Ok(bundle) = &result {
                self.returned_urls
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner)
                    .push(bundle.websocket_url.clone());
            }
            result
        })
    }
}

fn test_bundle(websocket_url: String) -> Result<NoiseRendezvousConnectBundle> {
    Ok(NoiseRendezvousConnectBundle {
        websocket_url,
        environment_id: "environment".to_string(),
        executor_registration_id: "registration".to_string(),
        executor_public_key: NoiseChannelIdentity::generate()?.public_key(),
        harness_key_authorization: "authorization".to_string(),
    })
}

fn registry_error(status: http::StatusCode, code: &str) -> ExecServerError {
    ExecServerError::EnvironmentRegistryHttp {
        status,
        code: Some(code.to_string()),
        message: "registry unavailable".to_string(),
    }
}

#[tokio::test]
async fn noise_handshake_uses_initialize_timeout() -> Result<()> {
    let listener = TcpListener::bind("127.0.0.1:0").await?;
    let websocket_url = format!("ws://{}", listener.local_addr()?);
    let server = tokio::spawn(async move {
        let (socket, _) = listener.accept().await?;
        let mut websocket = accept_async(socket).await?;
        // Drain the frames sent before the harness waits for the responder,
        // then verify that a timed-out readiness wait closes the socket.
        assert!(websocket.next().await.is_some());
        assert!(websocket.next().await.is_some());
        let closed =
            tokio::time::timeout(std::time::Duration::from_secs(1), websocket.next()).await?;
        assert!(
            matches!(closed, None | Some(Ok(Message::Close(_))) | Some(Err(_))),
            "timed-out Noise handshake must close its websocket"
        );
        anyhow::Ok(())
    });
    let initialize_timeout = std::time::Duration::from_millis(1);
    let opened = ExecServerClient::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs {
        bundle: NoiseRendezvousConnectBundle {
            websocket_url,
            environment_id: "environment".to_string(),
            executor_registration_id: "registration".to_string(),
            executor_public_key: NoiseChannelIdentity::generate()?.public_key(),
            harness_key_authorization: "authorization".to_string(),
        },
        harness_identity: NoiseChannelIdentity::generate()?,
        client_name: "test".to_string(),
        connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
        initialize_timeout,
        resume_session_id: None,
        http_client_factory: codex_http_client::HttpClientFactory::new(
            codex_http_client::OutboundProxyPolicy::ReqwestDefault,
        ),
    })
    .await?;

    let error = ExecServerClient::finish_noise_rendezvous_connection(opened)
        .await
        .err()
        .expect("stalled Noise handshake must time out");
    assert!(matches!(
        error,
        ExecServerError::InitializeTimedOut { timeout } if timeout == initialize_timeout
    ));

    server.await??;
    Ok(())
}

#[tokio::test]
async fn deferred_initialize_timeout_reports_configured_budget() {
    let (client_stdin, server_reader) = duplex(1 << 20);
    let (server_writer, client_stdout) = duplex(1 << 20);
    let server = tokio::spawn(async move {
        let _server_writer = server_writer;
        let mut lines = BufReader::new(server_reader).lines();
        let line = lines
            .next_line()
            .await
            .expect("initialize read should succeed")
            .expect("initialize request should arrive");
        let request: JSONRPCMessage =
            serde_json::from_str(&line).expect("initialize request should parse");
        assert!(
            matches!(
                request,
                JSONRPCMessage::Request(ref request) if request.method == INITIALIZE_METHOD
            ),
            "expected initialize request, got {request:?}"
        );
        futures::future::pending::<()>().await;
    });
    let configured_timeout = std::time::Duration::from_secs(10);
    let error = ExecServerClient::connect_with_recovery_and_noise_context(
        JsonRpcConnection::from_stdio(
            client_stdout,
            client_stdin,
            "timeout-test-client".to_string(),
        ),
        ExecServerClientConnectOptions {
            client_name: "timeout-test-client".to_string(),
            initialize_timeout: std::time::Duration::from_millis(1),
            resume_session_id: None,
        },
        /*reconnect_strategy*/ None,
        NoiseInitializeContext {
            span: tracing::info_span!("codex.exec_server.request"),
            timeout_for_error: configured_timeout,
        },
    )
    .await
    .err()
    .expect("initialize RPC must time out");
    assert!(matches!(
        error,
        ExecServerError::InitializeTimedOut { timeout } if timeout == configured_timeout
    ));
    server.abort();
    let _ = server.await;
}

#[tokio::test(start_paused = true)]
async fn initial_noise_connection_bounds_offline_retries() -> Result<()> {
    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
    for _ in 0..=INITIAL_REGISTRY_MAX_RETRIES {
        sequence.push_error(registry_error(
            http::StatusCode::CONFLICT,
            "environment_offline",
        ));
    }
    let identity = NoiseChannelIdentity::generate()?;
    let started = tokio::time::Instant::now();
    let error = sequence
        .connect(&identity)
        .await
        .err()
        .expect("offline retries must end");

    assert!(crate::client::is_environment_offline_error(&error));
    let requests = sequence.requested_keys().len();
    assert!((4..=INITIAL_REGISTRY_MAX_RETRIES as usize + 1).contains(&requests));
    sequence.assert_requested_identity(&identity, requests);
    assert!(started.elapsed() <= INITIAL_REGISTRY_OPERATION_TIMEOUT);
    Ok(())
}

#[tokio::test(start_paused = true)]
async fn initial_noise_connection_bounds_a_stalled_retry_request() -> Result<()> {
    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
    sequence.push_error(registry_error(
        http::StatusCode::CONFLICT,
        "environment_offline",
    ));
    for _ in 0..INITIAL_REGISTRY_MAX_RETRIES {
        sequence.push_pending();
    }
    let identity = NoiseChannelIdentity::generate()?;
    let started = tokio::time::Instant::now();
    let error = sequence
        .connect(&identity)
        .await
        .err()
        .expect("stalled retry must time out");

    assert!(matches!(
        error,
        ExecServerError::EnvironmentRegistryRequest(error) if error.is_timeout()
    ));
    assert_eq!(started.elapsed(), INITIAL_REGISTRY_OPERATION_TIMEOUT);
    let requests = sequence.requested_keys().len();
    assert!((2..=3).contains(&requests));
    sequence.assert_requested_identity(&identity, requests);
    Ok(())
}

#[tokio::test(start_paused = true)]
async fn initial_noise_connection_bounds_a_stalled_initial_request() -> Result<()> {
    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
    for _ in 0..=INITIAL_REGISTRY_MAX_RETRIES {
        sequence.push_pending();
    }
    let identity = NoiseChannelIdentity::generate()?;
    let started = tokio::time::Instant::now();

    let error = sequence
        .connect(&identity)
        .await
        .err()
        .expect("stalled initial request must time out");

    assert!(matches!(
        error,
        ExecServerError::EnvironmentRegistryRequest(error) if error.is_timeout()
    ));
    assert_eq!(started.elapsed(), INITIAL_REGISTRY_OPERATION_TIMEOUT);
    let requests = sequence.requested_keys().len();
    assert!((2..=3).contains(&requests));
    sequence.assert_requested_identity(&identity, requests);
    Ok(())
}

#[tokio::test(start_paused = true)]
async fn initial_noise_connection_retries_a_stalled_initial_request() -> Result<()> {
    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
    sequence.push_pending();
    sequence.push_error(registry_error(http::StatusCode::FORBIDDEN, "forbidden"));
    let identity = NoiseChannelIdentity::generate()?;
    let started = tokio::time::Instant::now();

    let error = sequence
        .connect(&identity)
        .await
        .err()
        .expect("terminal response must stop the retry sequence");

    assert!(matches!(
        error,
        ExecServerError::EnvironmentRegistryHttp {
            status: http::StatusCode::FORBIDDEN,
            ..
        }
    ));
    assert!(started.elapsed() >= INITIAL_REGISTRY_REQUEST_TIMEOUT);
    assert!(started.elapsed() < INITIAL_REGISTRY_OPERATION_TIMEOUT);
    sequence.assert_requested_identity(&identity, /*requests*/ 2);
    Ok(())
}

#[tokio::test(start_paused = true)]
async fn initial_noise_connection_retries_transient_registry_statuses() -> Result<()> {
    for status in [
        http::StatusCode::REQUEST_TIMEOUT,
        http::StatusCode::TOO_MANY_REQUESTS,
        http::StatusCode::INTERNAL_SERVER_ERROR,
        http::StatusCode::BAD_GATEWAY,
        http::StatusCode::SERVICE_UNAVAILABLE,
    ] {
        let sequence = Arc::new(SequenceNoiseConnectProvider::default());
        sequence.push_error(registry_error(status, "temporarily_unavailable"));
        sequence.push_error(registry_error(http::StatusCode::FORBIDDEN, "forbidden"));
        let identity = NoiseChannelIdentity::generate()?;

        let error = sequence
            .connect(&identity)
            .await
            .err()
            .expect("terminal response must stop the retry sequence");

        assert!(matches!(
            error,
            ExecServerError::EnvironmentRegistryHttp {
                status: http::StatusCode::FORBIDDEN,
                ..
            }
        ));
        sequence.assert_requested_identity(&identity, /*requests*/ 2);
    }
    Ok(())
}

#[tokio::test(start_paused = true)]
async fn initial_noise_connection_retries_registry_request_timeouts() -> Result<()> {
    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
    sequence.push_error(ExecServerError::EnvironmentRegistryRequest(
        codex_http_client::RouteAwareRequestError::Timeout,
    ));
    sequence.push_error(registry_error(http::StatusCode::FORBIDDEN, "forbidden"));
    let identity = NoiseChannelIdentity::generate()?;

    let error = sequence
        .connect(&identity)
        .await
        .err()
        .expect("terminal response must stop the retry sequence");

    assert!(matches!(
        error,
        ExecServerError::EnvironmentRegistryHttp {
            status: http::StatusCode::FORBIDDEN,
            ..
        }
    ));
    sequence.assert_requested_identity(&identity, /*requests*/ 2);
    Ok(())
}

#[tokio::test(start_paused = true)]
async fn initial_noise_connection_does_not_retry_permanent_registry_errors() -> Result<()> {
    for (status, code) in [
        (http::StatusCode::UNAUTHORIZED, "unauthorized"),
        (http::StatusCode::FORBIDDEN, "forbidden"),
        (http::StatusCode::BAD_REQUEST, "bad_request"),
        (http::StatusCode::NOT_FOUND, "environment_not_found"),
        (http::StatusCode::CONFLICT, "registration_conflict"),
        (http::StatusCode::CONFLICT, "route_unavailable"),
    ] {
        // A terminal error must also stop a retry sequence already in progress.
        for initial_offline in [false, true] {
            let sequence = Arc::new(SequenceNoiseConnectProvider::default());
            if initial_offline {
                sequence.push_error(registry_error(
                    http::StatusCode::CONFLICT,
                    "environment_offline",
                ));
            }
            sequence.push_error(registry_error(status, code));
            let identity = NoiseChannelIdentity::generate()?;
            let error = sequence
                .connect(&identity)
                .await
                .err()
                .expect("other errors must propagate");
            assert!(
                matches!(error, ExecServerError::EnvironmentRegistryHttp { status: actual_status, code: Some(actual_code), .. } if actual_status == status && actual_code == code)
            );
            sequence.assert_requested_identity(&identity, 1 + usize::from(initial_offline));
        }
    }
    Ok(())
}

#[tokio::test(start_paused = true)]
async fn noise_session_resume_leaves_offline_retries_to_recovery() -> Result<()> {
    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
    sequence.push_error(registry_error(
        http::StatusCode::CONFLICT,
        "environment_offline",
    ));
    let identity = NoiseChannelIdentity::generate()?;
    let strategy = ExecServerReconnectStrategy::NoiseRendezvous {
        executor_public_key: NoiseChannelIdentity::generate()?.public_key(),
        provider: sequence.clone(),
        identity: identity.clone(),
        client_name: "test".to_string(),
        connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
        initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
        http_client_factory: codex_http_client::HttpClientFactory::new(
            codex_http_client::OutboundProxyPolicy::ReqwestDefault,
        ),
    };
    let started = tokio::time::Instant::now();
    let error = strategy
        .resume("session")
        .await
        .err()
        .expect("resume must return the offline error");
    assert!(crate::client::is_environment_offline_error(&error));
    assert_eq!(started.elapsed(), std::time::Duration::ZERO);
    sequence.assert_requested_identity(&identity, /*requests*/ 1);
    Ok(())
}

#[tokio::test]
async fn initial_noise_connection_refreshes_bundle_after_exhausting_initial_retries() -> Result<()>
{
    let unauthorized_listener = TcpListener::bind("127.0.0.1:0").await?;
    let unauthorized_url = format!("ws://{}", unauthorized_listener.local_addr()?);
    let unauthorized_server = tokio::spawn(async move {
        let (mut socket, _) = unauthorized_listener.accept().await?;
        let mut request = [0_u8; 4096];
        let _ = socket.read(&mut request).await?;
        socket
            .write_all(
                b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
            )
            .await?;
        socket.shutdown().await?;
        anyhow::Ok(())
    });
    let accepted_listener = TcpListener::bind("127.0.0.1:0").await?;
    let accepted_url = format!("ws://{}", accepted_listener.local_addr()?);
    let executor_identity = NoiseChannelIdentity::generate()?;
    let executor_public_key = executor_identity.public_key();
    let accepted_server = tokio::spawn(async move {
        let (socket, _) = accepted_listener.accept().await?;
        let mut websocket = accept_async(socket).await?;
        let Message::Binary(resume_payload) = websocket.next().await.unwrap()? else {
            anyhow::bail!("expected Noise relay resume frame");
        };
        let resume = decode_relay_message_frame(resume_payload.as_ref())?;
        assert_eq!(resume.validate()?, RelayFrameBodyKind::Resume);
        let Message::Binary(handshake_payload) = websocket.next().await.unwrap()? else {
            anyhow::bail!("expected Noise relay handshake frame");
        };
        let handshake = decode_relay_message_frame(handshake_payload.as_ref())?;
        let stream_id = handshake.stream_id.clone();
        let prologue = noise_channel_prologue("environment", "registration", &stream_id);
        let pending = PendingResponderHandshake::read_request(
            &executor_identity,
            &prologue,
            &handshake.into_handshake_payload()?,
        )?;
        let (_transport, response) = pending.complete()?;
        websocket
            .send(Message::Binary(
                encode_relay_message_frame(&RelayMessageFrame::handshake(stream_id, response))
                    .into(),
            ))
            .await?;
        anyhow::Ok(())
    });
    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
    let unauthorized_bundle = test_bundle(unauthorized_url.clone())?;
    let mut accepted_bundle = test_bundle(accepted_url.clone())?;
    accepted_bundle.executor_public_key = executor_public_key;
    sequence.push_response(async {
        tokio::time::pause();
        Err(registry_error(
            http::StatusCode::CONFLICT,
            "environment_offline",
        ))
    });
    for _ in 1..INITIAL_REGISTRY_MAX_RETRIES {
        sequence.push_error(registry_error(
            http::StatusCode::CONFLICT,
            "environment_offline",
        ));
    }
    sequence.push_response(async move {
        tokio::time::resume();
        Ok(unauthorized_bundle)
    });
    sequence.push_response(async {
        tokio::time::pause();
        Err(registry_error(
            http::StatusCode::CONFLICT,
            "environment_offline",
        ))
    });
    sequence.push_response(async move {
        tokio::time::resume();
        Ok(accepted_bundle)
    });
    let identity = NoiseChannelIdentity::generate()?;

    let _connection = sequence.connect(&identity).await?;

    assert_eq!(
        sequence.returned_urls(),
        vec![unauthorized_url, accepted_url]
    );
    sequence.assert_requested_identity(&identity, INITIAL_REGISTRY_MAX_RETRIES as usize + 3);
    unauthorized_server.await??;
    accepted_server.await??;
    Ok(())
}