codex/codex-rs/rmcp-client/tests/streamable_http_test_support.rs

//! Shared helpers for Streamable HTTP RMCP integration tests.
//!
//! This support module starts the test HTTP server, launches a real
//! `exec-server` when remote coverage is needed, and provides small helpers for
//! creating RMCP clients and asserting round-trip behavior.
//! HTTP server startup uses the address published by the child that owns the listener.

// This support module is included by multiple integration-test crates. Each
// crate uses a different subset of the helpers, so dead-code warnings would
// otherwise depend on which test file compiled the module.
#![allow(dead_code)]

use std::net::SocketAddr;
use std::path::Path;
use std::path::PathBuf;
use std::process::Stdio;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;

use anyhow::Context as _;
use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_exec_server::Environment;
use codex_exec_server::ExecServerClient;
use codex_exec_server::HttpClient;
use codex_exec_server::RemoteExecServerConnectArgs;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_rmcp_client::ElicitationAction;
use codex_rmcp_client::ElicitationResponse;
use codex_rmcp_client::RmcpClient;
use codex_utils_cargo_bin::CargoBinError;
use futures::FutureExt as _;
use pretty_assertions::assert_eq;
use rmcp::model::CallToolResult;
use rmcp::model::ClientCapabilities;
use rmcp::model::ElicitationCapability;
use rmcp::model::FormElicitationCapability;
use rmcp::model::Implementation;
use rmcp::model::InitializeRequestParams;
use rmcp::model::ProtocolVersion;
use serde_json::json;
use tempfile::TempDir;
use tokio::io::AsyncBufReadExt;
use tokio::io::BufReader;
use tokio::process::Child;
use tokio::process::Command;
use tokio::time::sleep;

const SESSION_POST_FAILURE_CONTROL_PATH: &str = "/test/control/session-post-failure";
const INITIALIZE_POST_FAILURE_CONTROL_PATH: &str = "/test/control/initialize-post-failure";
const INITIALIZED_NOTIFICATION_POST_FAILURE_CONTROL_PATH: &str =
    "/test/control/initialized-notification-post-failure";

fn streamable_http_server_bin() -> Result<PathBuf, CargoBinError> {
    codex_utils_cargo_bin::cargo_bin("test_streamable_http_server")
}

fn init_params() -> InitializeRequestParams {
    let mut capabilities = ClientCapabilities::default();
    capabilities.elicitation =
        Some(ElicitationCapability::new().with_form(FormElicitationCapability::new()));
    InitializeRequestParams::new(
        capabilities,
        Implementation::new("codex-test", "0.0.0-test").with_title("Codex rmcp recovery test"),
    )
    .with_protocol_version(ProtocolVersion::V_2025_06_18)
}

pub(crate) fn expected_echo_result(message: &str) -> CallToolResult {
    let mut result = CallToolResult::success(Vec::new());
    result.result_type = None;
    result.structured_content = Some(json!({
        "echo": format!("ECHOING: {message}"),
        "env": null,
    }));
    result
}

pub(crate) async fn create_client(base_url: &str) -> anyhow::Result<RmcpClient> {
    create_client_with_http_client(base_url, Environment::default_for_tests().get_http_client())
        .await
}

pub(crate) async fn create_client_with_http_client(
    base_url: &str,
    http_client: Arc<dyn HttpClient>,
) -> anyhow::Result<RmcpClient> {
    let client = RmcpClient::new_streamable_http_client(
        "test-streamable-http",
        &format!("{base_url}/mcp"),
        Some("test-bearer".to_string()),
        /*http_headers*/ None,
        /*env_http_headers*/ None,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
        http_client,
        /*auth_provider*/ None,
    )
    .await?;

    initialize_client(&client).await?;

    Ok(client)
}

pub(crate) async fn initialize_client(client: &RmcpClient) -> anyhow::Result<()> {
    initialize_client_with_timeout(client, Duration::from_secs(/*secs*/ 5)).await
}

pub(crate) async fn initialize_client_with_timeout(
    client: &RmcpClient,
    timeout: Duration,
) -> anyhow::Result<()> {
    client
        .initialize(
            init_params(),
            Some(timeout),
            Box::new(|_, _| {
                async {
                    Ok(ElicitationResponse {
                        action: ElicitationAction::Accept,
                        content: Some(json!({})),
                        meta: None,
                    })
                }
                .boxed()
            }),
        )
        .await?;
    Ok(())
}

/// Creates a Streamable HTTP RMCP client that sends traffic through the remote
/// runtime HTTP API.
pub(crate) async fn create_remote_client(
    base_url: &str,
    http_client: ExecServerClient,
) -> anyhow::Result<RmcpClient> {
    let client = RmcpClient::new_streamable_http_client(
        "test-streamable-http-remote",
        &format!("{base_url}/mcp"),
        Some("test-bearer".to_string()),
        /*http_headers*/ None,
        /*env_http_headers*/ None,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
        Arc::new(http_client),
        /*auth_provider*/ None,
    )
    .await?;

    client
        .initialize(
            init_params(),
            Some(Duration::from_secs(5)),
            Box::new(|_, _| {
                async {
                    Ok(ElicitationResponse {
                        action: ElicitationAction::Accept,
                        content: Some(json!({})),
                        meta: None,
                    })
                }
                .boxed()
            }),
        )
        .await?;

    Ok(client)
}

pub(crate) async fn call_echo_tool(
    client: &RmcpClient,
    message: &str,
) -> anyhow::Result<CallToolResult> {
    client
        .call_tool(
            "echo".to_string(),
            Some(json!({ "message": message })),
            /*meta*/ None,
            Some(Duration::from_secs(5)),
        )
        .await
}

fn control_client() -> anyhow::Result<codex_http_client::HttpClient> {
    Ok(codex_http_client::HttpClientBuilder::new().build_direct()?)
}

pub(crate) async fn arm_session_post_failure(
    base_url: &str,
    status: u16,
    remaining: usize,
    www_authenticate_headers: &[&str],
) -> anyhow::Result<()> {
    let response = control_client()?
        .post(format!("{base_url}{SESSION_POST_FAILURE_CONTROL_PATH}"))
        .json(&json!({
            "status": status,
            "remaining": remaining,
            "www_authenticate_headers": www_authenticate_headers,
        }))
        .send()
        .await?;

    assert_eq!(response.status(), http::StatusCode::NO_CONTENT);
    Ok(())
}

pub(crate) async fn arm_session_post_json_rpc_failure(
    base_url: &str,
    status: u16,
    remaining: usize,
) -> anyhow::Result<()> {
    let response = control_client()?
        .post(format!("{base_url}{SESSION_POST_FAILURE_CONTROL_PATH}"))
        .json(&json!({
            "status": status,
            "remaining": remaining,
            "content_type": "application/json",
            "body": json!({
                "jsonrpc": "2.0",
                "id": 1,
                "error": {
                    "code": -32000,
                    "message": "transient session failure",
                },
            }).to_string(),
        }))
        .send()
        .await?;

    assert_eq!(response.status(), http::StatusCode::NO_CONTENT);
    Ok(())
}

pub(crate) async fn arm_initialized_notification_post_json_rpc_failure(
    base_url: &str,
    status: u16,
    remaining: usize,
) -> anyhow::Result<()> {
    let response = control_client()?
        .post(format!(
            "{base_url}{INITIALIZED_NOTIFICATION_POST_FAILURE_CONTROL_PATH}"
        ))
        .json(&json!({
            "status": status,
            "remaining": remaining,
            "content_type": "application/json",
            "body": json!({
                "jsonrpc": "2.0",
                "id": 1,
                "error": {
                    "code": -32000,
                    "message": "transient session failure",
                },
            }).to_string(),
        }))
        .send()
        .await?;

    assert_eq!(response.status(), http::StatusCode::NO_CONTENT);
    Ok(())
}

pub(crate) async fn arm_initialize_post_failure(
    base_url: &str,
    status: u16,
    remaining: usize,
) -> anyhow::Result<()> {
    let response = control_client()?
        .post(format!("{base_url}{INITIALIZE_POST_FAILURE_CONTROL_PATH}"))
        .json(&json!({
            "status": status,
            "remaining": remaining,
        }))
        .send()
        .await?;

    assert_eq!(response.status(), http::StatusCode::NO_CONTENT);
    Ok(())
}

pub(crate) async fn arm_initialize_post_json_rpc_failure(
    base_url: &str,
    status: u16,
    remaining: usize,
) -> anyhow::Result<()> {
    let response = control_client()?
        .post(format!("{base_url}{INITIALIZE_POST_FAILURE_CONTROL_PATH}"))
        .json(&json!({
            "status": status,
            "remaining": remaining,
            "content_type": "application/json",
            "body": json!({
                "jsonrpc": "2.0",
                "id": 1,
                "error": {
                    "code": -32000,
                    "message": "transient initialize failure",
                },
            }).to_string(),
        }))
        .send()
        .await?;

    assert_eq!(response.status(), http::StatusCode::NO_CONTENT);
    Ok(())
}

pub(crate) async fn spawn_streamable_http_server() -> anyhow::Result<(Child, String)> {
    let startup_dir = tempfile::tempdir()?;
    let bound_addr_path = startup_dir.path().join("bound_addr");
    // Let the child reserve its own port; releasing a parent-owned port before
    // spawning can let another test claim it and satisfy the readiness probe.
    let mut child = Command::new(streamable_http_server_bin()?)
        .kill_on_drop(true)
        .env("MCP_STREAMABLE_HTTP_BIND_ADDR", "127.0.0.1:0")
        .env("MCP_STREAMABLE_HTTP_BOUND_ADDR_FILE", &bound_addr_path)
        .spawn()?;

    let address = wait_for_streamable_http_server(
        &mut child,
        &bound_addr_path,
        Duration::from_secs(/*secs*/ 5),
    )
    .await?;
    Ok((child, format!("http://{address}")))
}

/// Owns the exec-server process used by the remote-client integration test.
pub(crate) struct ExecServerProcess {
    _codex_home: TempDir,
    child: Child,
    pub(crate) client: ExecServerClient,
}

impl Drop for ExecServerProcess {
    /// Stops the local exec-server process best-effort when the test exits.
    fn drop(&mut self) {
        let _ = self.child.start_kill();
    }
}

/// Starts a local exec-server and connects an initialized `ExecServerClient`.
pub(crate) async fn spawn_exec_server() -> anyhow::Result<ExecServerProcess> {
    let codex_home = TempDir::new()?;
    let mut child = Command::new(codex_utils_cargo_bin::cargo_bin("codex")?)
        .args(["exec-server", "--listen", "ws://127.0.0.1:0"])
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::inherit())
        .kill_on_drop(true)
        .env("CODEX_HOME", codex_home.path())
        .spawn()?;

    let websocket_url = read_exec_server_listen_url(&mut child).await?;
    let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new(
        websocket_url,
        "rmcp-client-remote-http-test".to_string(),
        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
    ))
    .await?;

    Ok(ExecServerProcess {
        _codex_home: codex_home,
        child,
        client,
    })
}

/// Reads the websocket URL printed by `codex exec-server --listen`.
async fn read_exec_server_listen_url(child: &mut Child) -> anyhow::Result<String> {
    let stdout = child
        .stdout
        .take()
        .context("failed to capture exec-server stdout")?;
    let mut lines = BufReader::new(stdout).lines();
    let deadline = Instant::now() + Duration::from_secs(10);

    loop {
        let remaining = deadline.saturating_duration_since(Instant::now());
        if remaining.is_zero() {
            anyhow::bail!("timed out waiting for exec-server listen URL");
        }

        let line = tokio::time::timeout(remaining, lines.next_line())
            .await
            .context("timed out waiting for exec-server stdout")??
            .context("exec-server stdout closed before emitting listen URL")?;
        let listen_url = line.trim();
        if listen_url.starts_with("ws://") {
            return Ok(listen_url.to_string());
        }
    }
}

async fn wait_for_streamable_http_server(
    server_child: &mut Child,
    bound_addr_path: &Path,
    timeout: Duration,
) -> anyhow::Result<SocketAddr> {
    let deadline = Instant::now() + timeout;

    loop {
        if let Some(status) = server_child.try_wait()? {
            return Err(anyhow::anyhow!(
                "streamable HTTP server exited early with status {status}"
            ));
        }

        if Instant::now() >= deadline {
            return Err(anyhow::anyhow!(
                "timed out waiting for streamable HTTP server to publish its bound address"
            ));
        }

        match std::fs::read_to_string(bound_addr_path) {
            Ok(contents) => {
                // The child creates the file before writing the address into it.
                if let Ok(address) = contents.parse() {
                    return Ok(address);
                }
            }
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
            Err(error) => {
                return Err(error).context("failed to read streamable HTTP server bound address");
            }
        }

        sleep(Duration::from_millis(50)).await;
    }
}