File size: 11,892 Bytes

ea39c0e

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_utils_path_uri::PathUri;
use pretty_assertions::assert_eq;
use tokio::sync::mpsc;
use uuid::Uuid;

use super::ExecServerHandler;
use crate::ExecServerRuntimePaths;
use crate::ProcessId;
use crate::protocol::ExecParams;
use crate::protocol::InitializeParams;
use crate::protocol::ReadParams;
use crate::protocol::ReadResponse;
use crate::protocol::TerminateParams;
use crate::protocol::TerminateResponse;
use crate::rpc::RpcNotificationSender;
use crate::server::session_registry::SessionRegistry;

fn exec_params(process_id: &str) -> ExecParams {
    exec_params_with_argv(process_id, sleep_argv())
}

fn exec_params_with_argv(process_id: &str, argv: Vec<String>) -> ExecParams {
    ExecParams {
        metadata: Default::default(),
        process_id: ProcessId::from(process_id),
        argv,
        cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd"))
            .expect("cwd URI"),
        shell_snapshot: None,
        env_policy: None,
        env: inherited_path_env(),
        tty: false,
        pipe_stdin: false,
        arg0: None,
        sandbox: None,
        enforce_managed_network: false,
        managed_network: None,
        network_proxy: None,
    }
}

fn inherited_path_env() -> HashMap<String, String> {
    let mut env = HashMap::new();
    if let Some(path) = std::env::var_os("PATH") {
        env.insert("PATH".to_string(), path.to_string_lossy().into_owned());
    }
    env
}

fn sleep_argv() -> Vec<String> {
    shell_argv("sleep 0.1", "ping -n 2 127.0.0.1 >NUL")
}

fn shell_argv(unix_script: &str, windows_script: &str) -> Vec<String> {
    if cfg!(windows) {
        vec![
            windows_command_processor(),
            "/C".to_string(),
            windows_script.to_string(),
        ]
    } else {
        vec![
            "/bin/sh".to_string(),
            "-c".to_string(),
            unix_script.to_string(),
        ]
    }
}

fn windows_command_processor() -> String {
    std::env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".to_string())
}

fn test_runtime_paths() -> ExecServerRuntimePaths {
    ExecServerRuntimePaths::new(
        std::env::current_exe().expect("current exe"),
        /*codex_linux_sandbox_exe*/ None,
    )
    .expect("runtime paths")
}

fn test_http_client_factory() -> HttpClientFactory {
    HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
}

async fn initialized_handler() -> Arc<ExecServerHandler> {
    let (outgoing_tx, _outgoing_rx) = mpsc::channel(16);
    let registry = SessionRegistry::new(crate::ExecServerTelemetry::default());
    let handler = Arc::new(ExecServerHandler::new(
        registry,
        RpcNotificationSender::new(outgoing_tx),
        test_runtime_paths(),
        test_http_client_factory(),
    ));
    let initialize_response = handler
        .initialize(InitializeParams {
            client_name: "exec-server-test".to_string(),
            resume_session_id: None,
        })
        .await
        .expect("initialize");
    Uuid::parse_str(&initialize_response.session_id).expect("session id should be a UUID");
    handler.initialized().expect("initialized");
    handler
}

#[tokio::test]
async fn duplicate_process_ids_allow_only_one_successful_start() {
    let handler = initialized_handler().await;
    let first_handler = Arc::clone(&handler);
    let second_handler = Arc::clone(&handler);

    let (first, second) = tokio::join!(
        first_handler.exec(exec_params("proc-1"), /*launch_context*/ None),
        second_handler.exec(exec_params("proc-1"), /*launch_context*/ None),
    );

    let (successes, failures): (Vec<_>, Vec<_>) =
        [first, second].into_iter().partition(Result::is_ok);
    assert_eq!(successes.len(), 1);
    assert_eq!(failures.len(), 1);

    let error = failures
        .into_iter()
        .next()
        .expect("one failed request")
        .expect_err("expected duplicate process error");
    assert_eq!(error.code, -32600);
    assert_eq!(error.message, "process proc-1 already exists");

    tokio::time::sleep(Duration::from_millis(150)).await;
    handler.shutdown().await;
}

#[tokio::test]
async fn terminate_reports_false_after_process_exit() {
    let handler = initialized_handler().await;
    handler
        .exec(exec_params("proc-1"), /*launch_context*/ None)
        .await
        .expect("start process");

    let deadline = tokio::time::Instant::now() + Duration::from_secs(1);
    loop {
        let response = handler
            .terminate(TerminateParams {
                process_id: ProcessId::from("proc-1"),
            })
            .await
            .expect("terminate response");
        if response == (TerminateResponse { running: false }) {
            break;
        }
        assert!(
            tokio::time::Instant::now() < deadline,
            "process should have exited within 1s"
        );
        tokio::time::sleep(Duration::from_millis(25)).await;
    }

    handler.shutdown().await;
}

#[tokio::test]
async fn long_poll_read_fails_after_session_resume() {
    let (first_tx, _first_rx) = mpsc::channel(16);
    let registry = SessionRegistry::new(crate::ExecServerTelemetry::default());
    let first_handler = Arc::new(ExecServerHandler::new(
        Arc::clone(&registry),
        RpcNotificationSender::new(first_tx),
        test_runtime_paths(),
        test_http_client_factory(),
    ));
    let initialize_response = first_handler
        .initialize(InitializeParams {
            client_name: "exec-server-test".to_string(),
            resume_session_id: None,
        })
        .await
        .expect("initialize");
    first_handler.initialized().expect("initialized");

    // Keep the process quiet and alive so the pending read can only complete
    // after session resume, not because the process produced output or exited.
    first_handler
        .exec(
            exec_params_with_argv(
                "proc-long-poll",
                shell_argv("sleep 5", "ping -n 6 127.0.0.1 >NUL"),
            ),
            /*launch_context*/ None,
        )
        .await
        .expect("start process");

    let first_read_handler = Arc::clone(&first_handler);
    let read_task = tokio::spawn(async move {
        first_read_handler
            .exec_read(ReadParams {
                process_id: ProcessId::from("proc-long-poll"),
                after_seq: None,
                max_bytes: None,
                wait_ms: Some(500),
            })
            .await
    });

    tokio::time::sleep(Duration::from_millis(50)).await;
    first_handler.shutdown().await;

    let (second_tx, _second_rx) = mpsc::channel(16);
    let second_handler = Arc::new(ExecServerHandler::new(
        registry,
        RpcNotificationSender::new(second_tx),
        test_runtime_paths(),
        test_http_client_factory(),
    ));
    second_handler
        .initialize(InitializeParams {
            client_name: "exec-server-test".to_string(),
            resume_session_id: Some(initialize_response.session_id),
        })
        .await
        .expect("initialize second connection");
    second_handler
        .initialized()
        .expect("initialized second connection");

    let err = read_task
        .await
        .expect("read task should join")
        .expect_err("evicted long-poll read should fail");
    assert_eq!(err.code, -32600);
    assert_eq!(
        err.message,
        "session has been resumed by another connection"
    );

    second_handler.shutdown().await;
}

#[tokio::test]
async fn active_session_resume_is_rejected() {
    let (first_tx, _first_rx) = mpsc::channel(16);
    let registry = SessionRegistry::new(crate::ExecServerTelemetry::default());
    let first_handler = Arc::new(ExecServerHandler::new(
        Arc::clone(&registry),
        RpcNotificationSender::new(first_tx),
        test_runtime_paths(),
        test_http_client_factory(),
    ));
    let initialize_response = first_handler
        .initialize(InitializeParams {
            client_name: "exec-server-test".to_string(),
            resume_session_id: None,
        })
        .await
        .expect("initialize");

    let (second_tx, _second_rx) = mpsc::channel(16);
    let second_handler = Arc::new(ExecServerHandler::new(
        registry,
        RpcNotificationSender::new(second_tx),
        test_runtime_paths(),
        test_http_client_factory(),
    ));
    let err = second_handler
        .initialize(InitializeParams {
            client_name: "exec-server-test".to_string(),
            resume_session_id: Some(initialize_response.session_id.clone()),
        })
        .await
        .expect_err("active session resume should fail");

    assert_eq!(err.code, crate::rpc::SESSION_ALREADY_ATTACHED_ERROR_CODE);
    assert_eq!(
        err.message,
        format!(
            "session {} is already attached to another connection",
            initialize_response.session_id
        )
    );

    first_handler.shutdown().await;
}

#[tokio::test]
async fn output_and_exit_are_retained_after_notification_receiver_closes() {
    let (outgoing_tx, outgoing_rx) = mpsc::channel(16);
    let handler = Arc::new(ExecServerHandler::new(
        SessionRegistry::new(crate::ExecServerTelemetry::default()),
        RpcNotificationSender::new(outgoing_tx),
        test_runtime_paths(),
        test_http_client_factory(),
    ));
    handler
        .initialize(InitializeParams {
            client_name: "exec-server-test".to_string(),
            resume_session_id: None,
        })
        .await
        .expect("initialize");
    handler.initialized().expect("initialized");

    let process_id = ProcessId::from("proc-notification-fail");
    handler
        .exec(
            exec_params_with_argv(
                process_id.as_str(),
                shell_argv(
                    "sleep 0.05; printf 'first\\n'; sleep 0.05; printf 'second\\n'",
                    "echo first&& ping -n 2 127.0.0.1 >NUL&& echo second",
                ),
            ),
            /*launch_context*/ None,
        )
        .await
        .expect("start process");

    drop(outgoing_rx);

    let (output, exit_code) = read_process_until_closed(&handler, process_id.clone()).await;
    assert_eq!(output.replace("\r\n", "\n"), "first\nsecond\n");
    assert_eq!(exit_code, Some(0));

    tokio::time::sleep(Duration::from_millis(100)).await;
    handler
        .exec(
            exec_params(process_id.as_str()),
            /*launch_context*/ None,
        )
        .await
        .expect("process id should be reusable after exit retention");

    handler.shutdown().await;
}

async fn read_process_until_closed(
    handler: &ExecServerHandler,
    process_id: ProcessId,
) -> (String, Option<i32>) {
    let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
    let mut output = String::new();
    let mut exit_code = None;
    let mut after_seq = None;

    loop {
        let response: ReadResponse = handler
            .exec_read(ReadParams {
                process_id: process_id.clone(),
                after_seq,
                max_bytes: None,
                wait_ms: Some(500),
            })
            .await
            .expect("read process");

        for chunk in response.chunks {
            output.push_str(&String::from_utf8_lossy(&chunk.chunk.into_inner()));
            after_seq = Some(chunk.seq);
        }
        if response.exited {
            exit_code = response.exit_code;
        }
        if response.closed {
            return (output, exit_code);
        }
        after_seq = response.next_seq.checked_sub(1).or(after_seq);
        assert!(
            tokio::time::Instant::now() < deadline,
            "process should close within 5s"
        );
    }
}