Remove WebSocket transport from code-mode hosts (#40692)

## What changed

- Accept only `http://` and `https://` gRPC endpoints for app server
  `--code-mode-host` connections.
- Limit the standalone code-mode host to stdio and gRPC listeners.
- Remove the WebSocket session provider, dual-WebSocket negotiation, transport
  lanes, and related dependencies and tests.
- Keep coverage for shared and prewarmed gRPC hosts and reject WebSocket host
  URLs in CLI parsing tests.

GitOrigin-RevId: 8792604e79ffd54ce776074a4dc3f7d55e63752a

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

@@ -1,43 +1,12 @@
 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;
 
@@ -86,274 +55,6 @@ fn missing_host_error_preserves_utf8_boundaries_when_truncating_the_path() {
     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(