Support remote code-mode hosts in app-server (#35098)

## What changed

- Add `--code-mode-host ws://...` and `wss://...` support to `codex app-server`, gated by the `code_mode_host` feature. When omitted, app-server continues to start a local host.
- Share one remote WebSocket connection across the process's threads, using the configured HTTP client's proxy and TLS policy and preserving the existing framed host protocol.
- Reject invalid host URLs, bound WebSocket frame sizes, close connections cleanly, and return an error when a connection exceeds 1,024 pending delegate calls without disconnecting it.

## Testing

- Cover CLI validation, WebSocket protocol execution and shutdown, connection sharing across app-server threads, and delegate-call capacity recovery.

GitOrigin-RevId: 715e82d4d9db1e7e2f91b754a777dcab504e2ae4

codex-rs/code-mode/Cargo.toml

@@ -17,11 +17,14 @@ workspace = true
 
 [dependencies]
 codex-code-mode-protocol = { workspace = true }
+codex-http-client = { workspace = true }
 codex-protocol = { workspace = true }
+codex-websocket-client = { workspace = true }
 deno_core_icudata = { workspace = true }
 futures = { workspace = true }
 serde_json = { workspace = true }
-tokio = { workspace = true, features = ["io-util", "macros", "process", "rt", "sync", "time"] }
+tokio = { workspace = true, features = ["io-util", "macros", "net", "process", "rt", "sync", "time"] }
+tokio-tungstenite = { workspace = true }
 tokio-util = { workspace = true, features = ["rt"] }
 tracing = { workspace = true }
 v8 = { workspace = true }

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

@@ -10,6 +10,7 @@ pub(crate) type TaskFailureHandler = std::sync::Arc<dyn Fn(String) + Send + Sync
 pub use codex_code_mode_protocol::*;
 pub use remote_session::ProcessOwnedCodeModeSession;
 pub use remote_session::ProcessOwnedCodeModeSessionProvider;
+pub use remote_session::WebSocketCodeModeSessionProvider;
 pub use service::InProcessCodeModeSession;
 pub use service::InProcessCodeModeSessionProvider;
 pub use service::NoopCodeModeSessionDelegate;

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

@@ -18,6 +18,8 @@ use codex_code_mode_protocol::StartedCell;
 use codex_code_mode_protocol::WaitOutcome;
 use codex_code_mode_protocol::WaitRequest;
 use codex_code_mode_protocol::host::SessionId;
+use codex_http_client::HttpClientFactory;
+use codex_http_client::OutboundProxyPolicy;
 use tokio::sync::Semaphore;
 use tokio::sync::watch;
 
@@ -38,8 +40,13 @@ pub struct ProcessOwnedCodeModeSessionProvider {
     state: StdMutex<ProviderState>,
 }
 
+/// Creates code-mode sessions backed by one shared remote WebSocket connection.
+pub struct WebSocketCodeModeSessionProvider {
+    host: Arc<OwnedCodeModeHost>,
+}
+
 enum ProviderState {
-    OwnedProcess(Arc<OwnedProcessHost>),
+    OwnedProcess(Arc<OwnedCodeModeHost>),
     InProcess,
 }
 
@@ -47,12 +54,12 @@ impl ProcessOwnedCodeModeSessionProvider {
     pub fn with_host_program(host_program: PathBuf) -> Self {
         Self {
             state: StdMutex::new(ProviderState::OwnedProcess(Arc::new(
-                OwnedProcessHost::new(host_program),
+                OwnedCodeModeHost::new(host_program),
             ))),
         }
     }
 
-    fn process_host(&self) -> Option<Arc<OwnedProcessHost>> {
+    fn process_host(&self) -> Option<Arc<OwnedCodeModeHost>> {
         match &*self
             .state
             .lock()
@@ -96,27 +103,84 @@ impl CodeModeSessionProvider for ProcessOwnedCodeModeSessionProvider {
                 }
                 Err(error) => return Err(error.to_string()),
             }
-            let session = ProcessOwnedCodeModeSession::with_process_host(delegate, process_host);
-            session.connection().await?;
-            let session: Arc<dyn CodeModeSession> = Arc::new(session);
-            Ok(session)
+            create_host_session(delegate, process_host).await
         })
     }
 }
 
-struct OwnedProcessHost {
-    host_program: PathBuf,
+impl WebSocketCodeModeSessionProvider {
+    pub fn new(websocket_url: String) -> Self {
+        Self::with_http_client_factory(
+            websocket_url,
+            HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+        )
+    }
+
+    /// Creates a remote host using the application's effective proxy and TLS policy.
+    pub fn with_http_client_factory(
+        websocket_url: String,
+        http_client_factory: HttpClientFactory,
+    ) -> Self {
+        Self {
+            host: Arc::new(OwnedCodeModeHost::websocket(
+                websocket_url,
+                http_client_factory,
+            )),
+        }
+    }
+}
+
+impl CodeModeSessionProvider for WebSocketCodeModeSessionProvider {
+    fn create_session<'a>(
+        &'a self,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+    ) -> CodeModeSessionProviderFuture<'a> {
+        Box::pin(create_host_session(delegate, Arc::clone(&self.host)))
+    }
+}
+
+async fn create_host_session(
+    delegate: Arc<dyn CodeModeSessionDelegate>,
+    host: Arc<OwnedCodeModeHost>,
+) -> Result<Arc<dyn CodeModeSession>, String> {
+    let session = ProcessOwnedCodeModeSession::with_host(delegate, host);
+    session.connection().await?;
+    Ok(Arc::new(session))
+}
+
+enum HostEndpoint {
+    Process(PathBuf),
+    WebSocket {
+        websocket_url: String,
+        http_client_factory: HttpClientFactory,
+    },
+}
+
+struct OwnedCodeModeHost {
+    endpoint: HostEndpoint,
     connection: StdMutex<Option<Arc<Connection>>>,
-    spawn_permit: Semaphore,
+    connect_permit: Semaphore,
     next_session_id: AtomicU64,
 }
 
-impl OwnedProcessHost {
+impl OwnedCodeModeHost {
     fn new(host_program: PathBuf) -> Self {
         Self {
-            host_program,
+            endpoint: HostEndpoint::Process(host_program),
             connection: StdMutex::new(None),
-            spawn_permit: Semaphore::new(/*permits*/ 1),
+            connect_permit: Semaphore::new(/*permits*/ 1),
+            next_session_id: AtomicU64::new(1),
+        }
+    }
+
+    fn websocket(websocket_url: String, http_client_factory: HttpClientFactory) -> Self {
+        Self {
+            endpoint: HostEndpoint::WebSocket {
+                websocket_url,
+                http_client_factory,
+            },
+            connection: StdMutex::new(None),
+            connect_permit: Semaphore::new(/*permits*/ 1),
             next_session_id: AtomicU64::new(1),
         }
     }
@@ -126,13 +190,20 @@ impl OwnedProcessHost {
             return Ok(connection);
         }
 
-        let _spawn_permit = self.spawn_permit.acquire().await.map_err(|_| {
-            ConnectionError::Other("code-mode host spawn coordinator closed".into())
+        let _connect_permit = self.connect_permit.acquire().await.map_err(|_| {
+            ConnectionError::Other("code-mode host connection coordinator closed".into())
         })?;
         if let Some(connection) = self.live_connection() {
             return Ok(connection);
         }
-        let new_connection = Arc::new(Connection::spawn(&self.host_program).await?);
+        let new_connection = match &self.endpoint {
+            HostEndpoint::Process(host_program) => Connection::spawn(host_program).await?,
+            HostEndpoint::WebSocket {
+                websocket_url,
+                http_client_factory,
+            } => Connection::connect_websocket(websocket_url, http_client_factory).await?,
+        };
+        let new_connection = Arc::new(new_connection);
         *self
             .connection
             .lock()
@@ -177,7 +248,7 @@ struct SessionBinding {
 }
 
 struct SessionInner {
-    process_host: Arc<OwnedProcessHost>,
+    host: Arc<OwnedCodeModeHost>,
     delegate: Arc<dyn CodeModeSessionDelegate>,
     state: StdMutex<SessionState>,
     next_generation: AtomicU64,
@@ -186,26 +257,23 @@ struct SessionInner {
     retired_cleanups: StdMutex<Vec<SessionCleanup>>,
 }
 
-/// A logical code-mode session assigned to a process-owned host.
+/// A logical code-mode session assigned to a process or WebSocket host.
 pub struct ProcessOwnedCodeModeSession {
     inner: Arc<SessionInner>,
 }
 
 impl ProcessOwnedCodeModeSession {
     pub fn new() -> Self {
-        Self::with_process_host(
+        Self::with_host(
             Arc::new(NoopCodeModeSessionDelegate),
-            Arc::new(OwnedProcessHost::new(default_host_program())),
+            Arc::new(OwnedCodeModeHost::new(default_host_program())),
         )
     }
 
-    fn with_process_host(
-        delegate: Arc<dyn CodeModeSessionDelegate>,
-        process_host: Arc<OwnedProcessHost>,
-    ) -> Self {
+    fn with_host(delegate: Arc<dyn CodeModeSessionDelegate>, host: Arc<OwnedCodeModeHost>) -> Self {
         Self {
             inner: Arc::new(SessionInner {
-                process_host,
+                host,
                 delegate,
                 state: StdMutex::new(SessionState::New),
                 next_generation: AtomicU64::new(1),
@@ -255,7 +323,7 @@ impl SessionInner {
                     SessionState::New => {
                         let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
                         let remote = RemoteSession {
-                            id: self.process_host.allocate_session_id(),
+                            id: self.host.allocate_session_id(),
                             generation,
                         };
                         let (result_tx, result_rx) = watch::channel(None);
@@ -294,7 +362,7 @@ impl SessionInner {
         remote: RemoteSession,
         result_tx: watch::Sender<Option<Result<SessionBinding, String>>>,
     ) {
-        let result = match self.process_host.connection().await {
+        let result = match self.host.connection().await {
             Ok(connection) => {
                 let cleanup = connection
                     .open_session(remote.clone(), Arc::clone(&self.delegate))

codex-rs/code-mode/src/remote_session/connection.rs

@@ -21,9 +21,13 @@ use codex_code_mode_protocol::host::EncodedFrame;
 use codex_code_mode_protocol::host::FramedReader;
 use codex_code_mode_protocol::host::FramedWriter;
 use codex_code_mode_protocol::host::HostToClient;
+use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
 use codex_code_mode_protocol::host::ProtocolVersion;
 use codex_code_mode_protocol::host::RequestId;
 use codex_code_mode_protocol::host::SupportedProtocolVersions;
+use codex_http_client::HttpClientFactory;
+use codex_websocket_client::WebSocketConnector;
+use futures::StreamExt;
 use tokio::io::AsyncBufReadExt;
 use tokio::io::BufReader;
 use tokio::process::Child;
@@ -31,6 +35,8 @@ use tokio::process::Command;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
 use tokio::task::JoinHandle;
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
 use tokio_util::sync::CancellationToken;
 use tracing::debug;
 use tracing::warn;
@@ -42,12 +48,16 @@ use self::driver::DriverLifecycle;
 pub(super) use self::driver::RemoteSession;
 pub(super) use self::driver::SessionCleanup;
 use self::reader::drive_reader;
+use self::transport::ConnectionReader;
+use self::transport::ConnectionWriter;
 
 mod driver;
 mod reader;
+mod transport;
 
 const IPC_CHANNEL_CAPACITY: usize = 128;
 const HOST_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
+const MAX_WEBSOCKET_FRAME_BYTES: usize = MAX_FRAME_BYTES + std::mem::size_of::<u32>();
 
 pub(super) enum ConnectionError {
     Spawn {
@@ -96,7 +106,7 @@ struct CallerCancellation {
 }
 
 struct ConnectionSupervisor {
-    child: Child,
+    owner: ConnectionOwner,
     event_tx: mpsc::Sender<DriverEvent>,
     cancellation: CancellationToken,
     alive: Arc<AtomicBool>,
@@ -106,6 +116,11 @@ struct ConnectionSupervisor {
     writer_task: JoinHandle<Result<(), String>>,
 }
 
+enum ConnectionOwner {
+    Process(Box<Child>),
+    WebSocket,
+}
+
 impl CallerCancellation {
     fn new() -> Self {
         Self {
@@ -171,8 +186,60 @@ impl Connection {
             .stdout
             .take()
             .ok_or_else(|| ConnectionError::Other("spawned code-mode host has no stdout".into()))?;
-        let mut reader = FramedReader::new(stdout);
-        let mut writer = FramedWriter::new(stdin);
+
+        Self::establish(
+            ConnectionReader::Stdio(FramedReader::new(stdout)),
+            ConnectionWriter::Stdio(FramedWriter::new(stdin)),
+            ConnectionOwner::Process(Box::new(child)),
+        )
+        .await
+    }
+
+    pub(super) async fn connect_websocket(
+        websocket_url: &str,
+        http_client_factory: &HttpClientFactory,
+    ) -> Result<Self, ConnectionError> {
+        let request = websocket_url.into_client_request().map_err(|error| {
+            ConnectionError::Other(format!(
+                "failed to build code-mode host websocket request: {error}"
+            ))
+        })?;
+        let connector = WebSocketConnector::new(http_client_factory).map_err(|error| {
+            ConnectionError::Other(format!(
+                "failed to configure code-mode host websocket TLS: {error}"
+            ))
+        })?;
+        let websocket_config = WebSocketConfig::default()
+            .max_frame_size(Some(MAX_WEBSOCKET_FRAME_BYTES))
+            .max_message_size(Some(MAX_WEBSOCKET_FRAME_BYTES));
+        let (websocket, _) = tokio::time::timeout(
+            HOST_HANDSHAKE_TIMEOUT,
+            connector.connect(request, websocket_config),
+        )
+        .await
+        .map_err(|_| {
+            ConnectionError::Other("timed out connecting to the code-mode host websocket".into())
+        })?
+        .map_err(|error| {
+            ConnectionError::Other(format!(
+                "failed to connect to the code-mode host websocket: {error}"
+            ))
+        })?;
+        let (writer, reader) = websocket.split();
+
+        Self::establish(
+            ConnectionReader::WebSocket(reader),
+            ConnectionWriter::WebSocket(writer),
+            ConnectionOwner::WebSocket,
+        )
+        .await
+    }
+
+    async fn establish(
+        mut reader: ConnectionReader,
+        mut writer: ConnectionWriter,
+        mut owner: ConnectionOwner,
+    ) -> Result<Self, ConnectionError> {
         let handshake = async {
             let hello = ClientHello::new(
                 SupportedProtocolVersions::try_new([ProtocolVersion::V1])
@@ -186,7 +253,7 @@ impl Connection {
                 .await
                 .map_err(|err| format!("failed to write code-mode host hello: {err}"))?;
             match reader
-                .read::<HostToClient>()
+                .read()
                 .await
                 .map_err(|err| format!("failed to read code-mode host hello: {err}"))?
             {
@@ -207,14 +274,16 @@ impl Connection {
         let handshake_result = match tokio::time::timeout(HOST_HANDSHAKE_TIMEOUT, handshake).await {
             Ok(result) => result,
             Err(_) => {
-                kill_and_reap(&mut child).await;
+                let _ = writer.close().await;
+                owner.close().await;
                 return Err(ConnectionError::Other(
                     "timed out negotiating with the code-mode host".into(),
                 ));
             }
         };
         if let Err(err) = handshake_result {
-            kill_and_reap(&mut child).await;
+            let _ = writer.close().await;
+            owner.close().await;
             return Err(ConnectionError::Other(err));
         }
 
@@ -229,12 +298,21 @@ impl Connection {
         let writer_task = tokio::spawn(async move {
             loop {
                 tokio::select! {
-                    _ = writer_cancellation.cancelled() => return Ok(()),
+                    _ = writer_cancellation.cancelled() => {
+                        return writer
+                            .close()
+                            .await
+                            .map_err(|error| format!("failed to close code-mode host connection: {error}"));
+                    }
                     frame = outgoing_rx.recv() => {
                         let Some(frame) = frame else {
                             return Err("code-mode host outgoing stream closed".to_string());
                         };
-                        if let Err(err) = writer.write_frame(&frame).await {
+                        let result = tokio::select! {
+                            _ = writer_cancellation.cancelled() => return Ok(()),
+                            result = writer.write_frame(frame) => result,
+                        };
+                        if let Err(err) = result {
                             return Err(format!("failed to write code-mode host message: {err}"));
                         }
                     }
@@ -263,7 +341,7 @@ impl Connection {
         let driver_task = tokio::spawn(driver.run());
         tokio::spawn(
             ConnectionSupervisor {
-                child,
+                owner,
                 event_tx,
                 cancellation: cancellation.clone(),
                 alive: Arc::clone(&alive),
@@ -428,7 +506,7 @@ impl Drop for Connection {
 
 impl ConnectionSupervisor {
     async fn run(mut self) {
-        let mut child_exited = false;
+        let mut owner_exited = false;
         let reason = tokio::select! {
             biased;
             _ = self.cancellation.cancelled() => failure_message(&self.failure),
@@ -438,19 +516,35 @@ impl ConnectionSupervisor {
             },
             result = &mut self.reader_task => task_failure("reader", result),
             result = &mut self.writer_task => task_failure("writer", result),
-            result = self.child.wait() => {
-                child_exited = true;
-                match result {
-                    Ok(status) => format!("code-mode host exited with status {status}"),
-                    Err(err) => format!("failed waiting for code-mode host: {err}"),
-                }
+            reason = self.owner.wait() => {
+                owner_exited = true;
+                reason
             }
         };
         mark_connection_dead(&self.alive, &self.failure, reason.clone());
         let _ = self.event_tx.try_send(DriverEvent::Failed(reason));
         self.cancellation.cancel();
-        if !child_exited {
-            kill_and_reap(&mut self.child).await;
+        if !owner_exited {
+            self.owner.close().await;
+        }
+    }
+}
+
+impl ConnectionOwner {
+    async fn wait(&mut self) -> String {
+        match self {
+            Self::Process(child) => match child.wait().await {
+                Ok(status) => format!("code-mode host exited with status {status}"),
+                Err(error) => format!("failed waiting for code-mode host: {error}"),
+            },
+            Self::WebSocket => std::future::pending().await,
+        }
+    }
+
+    async fn close(&mut self) {
+        match self {
+            Self::Process(child) => kill_and_reap(child).await,
+            Self::WebSocket => {}
         }
     }
 }

codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs

@@ -13,6 +13,7 @@ use codex_code_mode_protocol::host::DelegateRequest;
 use codex_code_mode_protocol::host::DelegateRequestId;
 use codex_code_mode_protocol::host::DelegateResponse;
 use codex_code_mode_protocol::host::EncodedFrame;
+use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
 use codex_code_mode_protocol::host::SessionId;
 use codex_code_mode_protocol::host::WireCellId;
 use codex_code_mode_protocol::host::WireResult;
@@ -65,6 +66,11 @@ enum DelegateTask {
     },
 }
 
+enum DelegateStartError {
+    Duplicate(DelegateRequestId),
+    CapacityExceeded,
+}
+
 pub(super) struct DelegateEffects {
     pub(super) response: Option<(DelegateRequestId, Result<DelegateResponse, String>)>,
     pub(super) closed_cells: Vec<CellOwner>,
@@ -102,16 +108,19 @@ impl DelegateRuntime {
         }
     }
 
-    pub(super) fn start(
+    fn start(
         &mut self,
         id: DelegateRequestId,
         target: DelegateTarget,
         request: DelegateRequest,
-    ) -> Result<(), String> {
+    ) -> Result<(), DelegateStartError> {
         if self.calls.contains_key(&id) || self.seen_requests.contains(&id) {
-            return Err(format!("duplicate code-mode delegate request ID {id:?}"));
+            return Err(DelegateStartError::Duplicate(id));
         }
         self.remember_request(id);
+        if self.calls.len() >= MAX_PENDING_DELEGATE_CALLS {
+            return Err(DelegateStartError::CapacityExceeded);
+        }
         let cancellation = CancellationToken::new();
         let task_request = match request {
             DelegateRequest::InvokeTool { invocation } => {
@@ -257,11 +266,19 @@ impl ConnectionDriver {
                 return false;
             }
         };
-        if let Err(err) = self.delegates.start(id, target, request) {
-            self.fail(err);
-            return false;
+        match self.delegates.start(id, target, request) {
+            Ok(()) => true,
+            Err(DelegateStartError::Duplicate(id)) => {
+                self.fail(format!("duplicate code-mode delegate request ID {id:?}"));
+                false
+            }
+            Err(DelegateStartError::CapacityExceeded) => self.send_delegate_response(
+                id,
+                Err(format!(
+                    "code-mode host exceeded the limit of {MAX_PENDING_DELEGATE_CALLS} pending delegate calls"
+                )),
+            ),
         }
-        true
     }
 
     pub(super) fn complete_delegate(

codex-rs/code-mode/src/remote_session/connection/driver_tests.rs

@@ -12,10 +12,13 @@ use codex_code_mode_protocol::ExecuteRequest;
 use codex_code_mode_protocol::NotificationFuture;
 use codex_code_mode_protocol::ToolInvocationFuture;
 use codex_code_mode_protocol::WaitRequest;
+use codex_code_mode_protocol::host::ClientToHost;
 use codex_code_mode_protocol::host::DelegateRequest;
 use codex_code_mode_protocol::host::DelegateRequestId;
+use codex_code_mode_protocol::host::EncodedFrame;
 use codex_code_mode_protocol::host::HostResponse;
 use codex_code_mode_protocol::host::HostToClient;
+use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
 use codex_code_mode_protocol::host::RequestId;
 use codex_code_mode_protocol::host::SessionId;
 use codex_code_mode_protocol::host::WireNestedToolCall;
@@ -462,6 +465,73 @@ async fn delegate_cancel_is_best_effort_and_sends_no_late_response() {
     assert_eq!(delegate.notifications.load(Ordering::Relaxed), 0);
 }
 
+#[tokio::test]
+async fn delegate_limit_returns_an_error_without_disconnecting() {
+    let mut harness = DriverHarness::start();
+    let session = remote_session();
+    let delegate = Arc::new(RecordingDelegate::default());
+    harness.open(session.clone(), delegate.clone()).await;
+    let _started = harness
+        .start_cell(session.clone(), /*request_id*/ 2, "1")
+        .await;
+
+    for value in 1..=MAX_PENDING_DELEGATE_CALLS {
+        harness
+            .start_tool_delegate(&session, DelegateRequestId::new(value as i64))
+            .await;
+    }
+
+    let overflow_id = DelegateRequestId::new(MAX_PENDING_DELEGATE_CALLS as i64 + 1);
+    harness.start_tool_delegate(&session, overflow_id).await;
+    let response = tokio::time::timeout(Duration::from_secs(5), harness.outgoing_rx.recv())
+        .await
+        .expect("delegate overflow response timeout")
+        .expect("delegate overflow response frame");
+
+    assert_eq!(
+        EncodedFrame::decode_framed::<ClientToHost>(&response.into_framed_bytes())
+            .expect("decode delegate overflow response"),
+        ClientToHost::DelegateResponse {
+            id: overflow_id,
+            result: WireResult::Err {
+                message: format!(
+                    "code-mode host exceeded the limit of {MAX_PENDING_DELEGATE_CALLS} pending delegate calls"
+                ),
+            },
+        }
+    );
+    assert!(harness.alive.load(Ordering::Acquire));
+
+    harness
+        .event_tx
+        .send(DriverEvent::HostMessage(
+            HostToClient::CancelDelegateRequest {
+                id: DelegateRequestId::new(/*value*/ 1),
+            },
+        ))
+        .await
+        .expect("cancel pending delegate");
+    harness
+        .start_tool_delegate(
+            &session,
+            DelegateRequestId::new(MAX_PENDING_DELEGATE_CALLS as i64 + 2),
+        )
+        .await;
+
+    tokio::time::timeout(Duration::from_secs(5), async {
+        while delegate.invocations.load(Ordering::Relaxed) <= MAX_PENDING_DELEGATE_CALLS {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("delegate capacity should be available after cancellation");
+    assert_eq!(
+        delegate.invocations.load(Ordering::Relaxed),
+        MAX_PENDING_DELEGATE_CALLS + 1
+    );
+    assert!(harness.alive.load(Ordering::Acquire));
+}
+
 #[tokio::test]
 async fn terminate_closes_cell_without_waiting_for_delegate_cleanup() {
     let mut harness = DriverHarness::start();

codex-rs/code-mode/src/remote_session/connection/reader.rs

@@ -1,20 +1,18 @@
-use codex_code_mode_protocol::host::FramedReader;
-use codex_code_mode_protocol::host::HostToClient;
-use tokio::process::ChildStdout;
 use tokio::sync::mpsc;
 use tokio_util::sync::CancellationToken;
 
 use super::driver::DriverEvent;
+use super::transport::ConnectionReader;
 
 pub(super) async fn drive_reader(
-    mut reader: FramedReader<ChildStdout>,
+    mut reader: ConnectionReader,
     events: mpsc::Sender<DriverEvent>,
     cancellation: CancellationToken,
 ) -> Result<(), String> {
     loop {
         let message = tokio::select! {
             _ = cancellation.cancelled() => return Ok(()),
-            result = reader.read::<HostToClient>() => result,
+            result = reader.read() => result,
         };
         let message = match message {
             Ok(Some(message)) => message,

codex-rs/code-mode/src/remote_session/connection/transport.rs

@@ -0,0 +1,103 @@
+use std::io;
+use std::time::Duration;
+
+use codex_code_mode_protocol::host::ClientToHost;
+use codex_code_mode_protocol::host::EncodedFrame;
+use codex_code_mode_protocol::host::FramedReader;
+use codex_code_mode_protocol::host::FramedWriter;
+use codex_code_mode_protocol::host::HostToClient;
+use codex_websocket_client::WebSocketConnection;
+use futures::SinkExt;
+use futures::StreamExt;
+use futures::stream::SplitSink;
+use futures::stream::SplitStream;
+use tokio::process::ChildStdin;
+use tokio::process::ChildStdout;
+use tokio_tungstenite::tungstenite::Message;
+
+const WEBSOCKET_CLOSE_TIMEOUT: Duration = Duration::from_secs(5);
+
+pub(super) enum ConnectionReader {
+    Stdio(FramedReader<ChildStdout>),
+    WebSocket(SplitStream<WebSocketConnection>),
+}
+
+pub(super) enum ConnectionWriter {
+    Stdio(FramedWriter<ChildStdin>),
+    WebSocket(SplitSink<WebSocketConnection, Message>),
+}
+
+impl ConnectionReader {
+    pub(super) async fn read(&mut self) -> io::Result<Option<HostToClient>> {
+        match self {
+            Self::Stdio(reader) => reader.read().await,
+            Self::WebSocket(reader) => loop {
+                match reader.next().await {
+                    Some(Ok(Message::Binary(frame))) => {
+                        return EncodedFrame::decode_framed(&frame).map(Some);
+                    }
+                    Some(Ok(Message::Ping(_) | Message::Pong(_))) => {}
+                    Some(Ok(Message::Close(_))) | None => return Ok(None),
+                    Some(Ok(Message::Text(_))) => {
+                        return Err(io::Error::new(
+                            io::ErrorKind::InvalidData,
+                            "code-mode host websocket messages must be binary framed messages",
+                        ));
+                    }
+                    Some(Ok(Message::Frame(_))) => {
+                        return Err(io::Error::new(
+                            io::ErrorKind::InvalidData,
+                            "code-mode host websocket returned an unexpected raw frame",
+                        ));
+                    }
+                    Some(Err(error)) => {
+                        return Err(io::Error::other(format!(
+                            "failed to read code-mode host websocket message: {error}"
+                        )));
+                    }
+                }
+            },
+        }
+    }
+}
+
+impl ConnectionWriter {
+    pub(super) async fn write(&mut self, message: &ClientToHost) -> io::Result<()> {
+        self.write_frame(EncodedFrame::encode(message)?).await
+    }
+
+    pub(super) async fn write_frame(&mut self, frame: EncodedFrame) -> io::Result<()> {
+        match self {
+            Self::Stdio(writer) => writer.write_frame(&frame).await,
+            Self::WebSocket(writer) => writer
+                .send(Message::Binary(frame.into_framed_bytes().into()))
+                .await
+                .map_err(|error| {
+                    io::Error::other(format!(
+                        "failed to write code-mode host websocket message: {error}"
+                    ))
+                }),
+        }
+    }
+
+    pub(super) async fn close(&mut self) -> io::Result<()> {
+        match self {
+            Self::Stdio(_) => Ok(()),
+            Self::WebSocket(writer) => {
+                tokio::time::timeout(WEBSOCKET_CLOSE_TIMEOUT, writer.close())
+                    .await
+                    .map_err(|_| {
+                        io::Error::new(
+                            io::ErrorKind::TimedOut,
+                            "timed out closing code-mode host websocket connection",
+                        )
+                    })?
+                    .map_err(|error| {
+                        io::Error::other(format!(
+                            "failed to close code-mode host websocket connection: {error}"
+                        ))
+                    })
+            }
+        }
+    }
+}

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

@@ -1,15 +1,37 @@
 use std::io;
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::time::Duration;
 
 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::CapabilitySet;
+use codex_code_mode_protocol::host::ClientToHost;
+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::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::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::resolve_host_program;
 use crate::NoopCodeModeSessionDelegate;
 
@@ -107,6 +129,150 @@ async fn provider_falls_back_to_in_process_session_when_host_is_missing() {
     );
 }
 
+#[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(_) => 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 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 shutdown_before_open_does_not_spawn_the_host() {
     let session = ProcessOwnedCodeModeSession::new();