codex/codex-rs/code-mode/src/remote_session_tests.rs

use std::io;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionProvider;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::FunctionCallOutputContentItem;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientToHost;
use codex_code_mode_protocol::host::DUAL_WEBSOCKET_CAPABILITY;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::HostHello;
use codex_code_mode_protocol::host::HostRequest;
use codex_code_mode_protocol::host::HostResponse;
use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::ProtocolVersion;
use codex_code_mode_protocol::host::SESSION_RESOURCE_LIMITS_CAPABILITY;
use codex_code_mode_protocol::host::WireCellId;
use codex_code_mode_protocol::host::WireContentItem;
use codex_code_mode_protocol::host::WireResult;
use codex_code_mode_protocol::host::WireRuntimeResponse;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use futures::SinkExt;
use futures::StreamExt;
use pretty_assertions::assert_eq;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpListener;
use tokio::time::timeout;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::tungstenite::Message;

use super::ProcessOwnedCodeModeSession;
use super::ProcessOwnedCodeModeSessionProvider;
use super::WebSocketCodeModeSessionProvider;
use super::connection::ConnectionError;
use crate::NoopCodeModeSessionDelegate;

#[test]
fn provider_reuses_its_live_process_host() {
    let provider = ProcessOwnedCodeModeSessionProvider::default();

    let first = provider.process_host();
    let second = provider.process_host();

    assert!(Arc::ptr_eq(&first, &second));
}

#[test]
fn missing_host_error_limits_the_displayed_path_to_512_bytes() {
    let executable = "codex-code-mode-host-does-not-exist";
    let host_program = format!("{}{executable}", "missing-directory/".repeat(/*n*/ 64));
    let expected_suffix = &host_program[host_program.len() - (512 - "...".len())..];
    let error = ConnectionError::Spawn {
        host_program: PathBuf::from(&host_program),
        error: io::Error::new(io::ErrorKind::NotFound, "host unavailable"),
    };

    assert_eq!(
        error.to_string(),
        format!("failed to spawn code-mode host ...{expected_suffix}: host unavailable")
    );
}

#[test]
fn missing_host_error_preserves_utf8_boundaries_when_truncating_the_path() {
    let executable = "codex-code-mode-host-does-not-exist";
    let host_program = format!("{}{executable}", "🦀".repeat(/*n*/ 256));
    let error = ConnectionError::Spawn {
        host_program: PathBuf::from(host_program),
        error: io::Error::new(io::ErrorKind::NotFound, "host unavailable"),
    }
    .to_string();
    let displayed_path = error
        .strip_prefix("failed to spawn code-mode host ")
        .and_then(|message| message.strip_suffix(": host unavailable"))
        .expect("missing-host error should contain the displayed host path");

    assert!(displayed_path.starts_with("..."));
    assert!(displayed_path.ends_with(executable));
    assert!(displayed_path.len() <= 512);
}

#[tokio::test]
async fn websocket_provider_executes_over_shared_connector() {
    let listener = TcpListener::bind("127.0.0.1:0")
        .await
        .expect("websocket test listener should bind");
    let websocket_url = format!(
        "ws://{}",
        listener
            .local_addr()
            .expect("websocket test listener should have an address")
    );
    let server = tokio::spawn(async move {
        let (stream, _) = listener
            .accept()
            .await
            .expect("websocket test host should accept a connection");
        let mut websocket = accept_async(stream)
            .await
            .expect("websocket test host should complete the HTTP handshake");

        while let Some(message) = websocket.next().await {
            let message = message.expect("websocket test host should receive a valid message");
            let frame = match message {
                Message::Binary(frame) => frame,
                Message::Ping(_) | Message::Pong(_) => continue,
                Message::Close(_) => break,
                Message::Text(_) | Message::Frame(_) => {
                    panic!("websocket test host received an unexpected message: {message:?}");
                }
            };
            let request = EncodedFrame::decode_framed::<ClientToHost>(&frame)
                .expect("websocket test host should decode a framed protocol message");
            let responses = match request {
                ClientToHost::ClientHello(hello) => {
                    let capability = Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY)
                        .expect("session-limit capability");
                    assert!(hello.optional_capabilities().contains(&capability));
                    assert_eq!(hello.required_capabilities(), &CapabilitySet::empty());
                    vec![HostToClient::HostHello(HostHello::new(
                        ProtocolVersion::V1,
                        CapabilitySet::empty(),
                    ))]
                }
                ClientToHost::Request {
                    id,
                    request: HostRequest::OpenSession { session_id, .. },
                } => vec![HostToClient::Response {
                    id,
                    result: WireResult::Ok {
                        value: HostResponse::SessionReady { session_id },
                    },
                }],
                ClientToHost::Request {
                    id,
                    request: HostRequest::Execute { request, .. },
                } => {
                    assert_eq!(request.source, "text('shared connector')");
                    let cell_id = WireCellId::new("1");
                    vec![
                        HostToClient::Response {
                            id,
                            result: WireResult::Ok {
                                value: HostResponse::ExecutionStarted {
                                    cell_id: cell_id.clone(),
                                },
                            },
                        },
                        HostToClient::InitialResponse {
                            id,
                            result: WireResult::Ok {
                                value: WireRuntimeResponse::Result {
                                    cell_id,
                                    content_items: vec![WireContentItem::InputText {
                                        text: "shared connector".to_string(),
                                    }],
                                    error_text: None,
                                },
                            },
                        },
                    ]
                }
                ClientToHost::Request {
                    id,
                    request: HostRequest::ShutdownSession { session_id },
                } => vec![HostToClient::Response {
                    id,
                    result: WireResult::Ok {
                        value: HostResponse::SessionClosed { session_id },
                    },
                }],
                request => {
                    panic!("websocket test host received an unexpected request: {request:?}")
                }
            };

            for response in responses {
                let frame = EncodedFrame::encode(&response)
                    .expect("websocket test host should encode a framed response");
                websocket
                    .send(Message::Binary(frame.into_framed_bytes().into()))
                    .await
                    .expect("websocket test host should send its response");
            }
        }
    });

    let provider = WebSocketCodeModeSessionProvider::with_http_client_factory(
        websocket_url,
        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
    );
    let session = provider
        .create_session(Arc::new(NoopCodeModeSessionDelegate))
        .await
        .expect("shared websocket connector should open a code-mode session");
    let error = provider
        .create_session_with_limits(
            Arc::new(NoopCodeModeSessionDelegate),
            CodeModeSessionCellExecutionLimits {
                max_yield_time_ms: Some(250),
                max_heap_size_bytes: None,
            },
        )
        .await
        .err()
        .expect("legacy host should reject a limited session");
    assert_eq!(
        error,
        format!(
            "code-mode host does not support session resource limits: missing `{SESSION_RESOURCE_LIMITS_CAPABILITY}` capability"
        )
    );
    let second_session = provider
        .create_session(Arc::new(NoopCodeModeSessionDelegate))
        .await
        .expect("rejecting limited sessions should preserve the shared legacy-host connection");
    second_session
        .shutdown()
        .await
        .expect("second unlimited session should shut down");
    drop(second_session);
    let response = session
        .execute(ExecuteRequest {
            tool_call_id: "shared-websocket".to_string(),
            enabled_tools: Vec::new(),
            source: "text('shared connector')".to_string(),
            yield_time_ms: None,
            max_output_tokens: None,
        })
        .await
        .expect("shared websocket connector should start a cell")
        .initial_response()
        .await
        .expect("shared websocket connector should return a cell result");

    assert_eq!(
        response,
        RuntimeResponse::Result {
            cell_id: codex_code_mode_protocol::CellId::new("1".to_string()),
            content_items: vec![FunctionCallOutputContentItem::InputText {
                text: "shared connector".to_string(),
            }],
            error_text: None,
        }
    );
    session
        .shutdown()
        .await
        .expect("shared websocket connector should shut down its session");
    drop(session);
    drop(provider);
    timeout(Duration::from_secs(5), server)
        .await
        .expect("websocket test host should disconnect promptly")
        .expect("websocket test host task should succeed");
}

#[tokio::test]
async fn websocket_provider_fails_when_a_negotiated_bulk_connection_is_unavailable() {
    let listener = TcpListener::bind("127.0.0.1:0")
        .await
        .expect("websocket test listener should bind");
    let websocket_url = format!(
        "ws://{}/?access_token=shared-token",
        listener
            .local_addr()
            .expect("websocket test listener should have an address")
    );
    let server = tokio::spawn(async move {
        let (stream, _) = listener
            .accept()
            .await
            .expect("websocket test host should accept the first connection");
        let mut control = accept_async(stream)
            .await
            .expect("first control websocket should connect");
        let frame = control
            .next()
            .await
            .expect("first client hello")
            .expect("first websocket frame")
            .into_data();
        let ClientToHost::ClientHello(hello) =
            EncodedFrame::decode_framed(&frame).expect("decode first client hello")
        else {
            panic!("expected first client hello");
        };
        let capability =
            Capability::new(DUAL_WEBSOCKET_CAPABILITY).expect("dual websocket capability");
        assert!(hello.optional_capabilities().contains(&capability));
        let session_limits_capability =
            Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY).expect("session-limit capability");
        assert!(
            hello
                .optional_capabilities()
                .contains(&session_limits_capability)
        );
        let hello = HostToClient::HostHello(
            HostHello::new(
                ProtocolVersion::V1,
                CapabilitySet::try_new([capability.clone()]).expect("host capabilities"),
            )
            .with_bulk_connection_token("fallback-token".to_string()),
        );
        let frame = EncodedFrame::encode(&hello).expect("encode dual host hello");
        control
            .send(Message::Binary(frame.into_framed_bytes().into()))
            .await
            .expect("send dual host hello");

        let (mut bulk, _) = listener
            .accept()
            .await
            .expect("websocket test host should accept the bulk connection");
        let mut request = [0_u8; 1024];
        let request_len = bulk
            .read(&mut request)
            .await
            .expect("read bulk websocket handshake");
        let request = std::str::from_utf8(&request[..request_len]).expect("bulk HTTP request");
        assert!(
            request.starts_with("GET /bulk/fallback-token?access_token=shared-token HTTP/1.1\r\n")
        );
        bulk.write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
            .await
            .expect("reject unavailable bulk websocket");
        drop(bulk);
        drop(control);
    });

    let provider = WebSocketCodeModeSessionProvider::new(websocket_url);
    let error = match provider
        .create_session(Arc::new(NoopCodeModeSessionDelegate))
        .await
    {
        Ok(_) => panic!("provider should reject an unavailable negotiated bulk websocket"),
        Err(error) => error,
    };
    assert!(
        error.contains("404"),
        "unexpected negotiated bulk websocket error: {error}"
    );
    drop(provider);
    timeout(Duration::from_secs(5), server)
        .await
        .expect("negotiated bulk websocket test host should disconnect promptly")
        .expect("negotiated bulk websocket test host task should succeed");
}

#[tokio::test]
async fn provider_returns_missing_host_error() {
    let provider = ProcessOwnedCodeModeSessionProvider::with_host_program(
        "codex-code-mode-host-does-not-exist".into(),
    );

    let error = provider
        .create_session(Arc::new(NoopCodeModeSessionDelegate))
        .await
        .err()
        .expect("missing host should fail");

    assert!(error.contains("failed to spawn code-mode host codex-code-mode-host-does-not-exist"));
}

#[tokio::test]
async fn shutdown_before_open_does_not_spawn_the_host() {
    let session = ProcessOwnedCodeModeSession::new();

    session.shutdown().await.expect("shutdown session");
    let error = session
        .execute(codex_code_mode_protocol::ExecuteRequest {
            tool_call_id: "call-1".to_string(),
            enabled_tools: Vec::new(),
            source: "text('unreachable')".to_string(),
            yield_time_ms: None,
            max_output_tokens: None,
        })
        .await
        .err()
        .expect("shutdown session should reject execution");

    assert_eq!(error, "code mode session is shutting down");
}