From f61b51ddd924643514b33234816a8a2772b1aec7 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Fri, 24 Jul 2026 04:32:21 +0000 Subject: [PATCH] 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/Cargo.lock | 4 + codex-rs/app-server/Cargo.toml | 3 +- codex-rs/app-server/README.md | 2 + codex-rs/app-server/src/code_mode_host.rs | 48 +++++ .../app-server/src/code_mode_host_tests.rs | 51 ++++++ codex-rs/app-server/src/in_process.rs | 1 + codex-rs/app-server/src/lib.rs | 29 ++- codex-rs/app-server/src/main.rs | 10 +- codex-rs/app-server/src/main_tests.rs | 36 ++++ codex-rs/app-server/src/message_processor.rs | 11 +- .../src/message_processor_tracing_tests.rs | 1 + .../tests/suite/v2/code_mode_host.rs | 137 +++++++++++++++ codex-rs/app-server/tests/suite/v2/mod.rs | 1 + .../tests/suite/v2/remote_control.rs | 1 + codex-rs/cli/src/main.rs | 49 ++++++ codex-rs/code-mode-host/src/peer.rs | 2 +- codex-rs/code-mode-protocol/src/host/mod.rs | 3 + codex-rs/code-mode/Cargo.toml | 5 +- codex-rs/code-mode/src/lib.rs | 1 + codex-rs/code-mode/src/remote_session.rs | 122 ++++++++++--- .../src/remote_session/connection.rs | 130 ++++++++++++-- .../connection/driver/delegate_runtime.rs | 31 +++- .../remote_session/connection/driver_tests.rs | 70 ++++++++ .../src/remote_session/connection/reader.rs | 8 +- .../remote_session/connection/transport.rs | 103 +++++++++++ .../code-mode/src/remote_session_tests.rs | 166 ++++++++++++++++++ codex-rs/core/src/thread_manager.rs | 12 ++ codex-rs/core/src/thread_manager_tests.rs | 5 +- 28 files changed, 977 insertions(+), 65 deletions(-) create mode 100644 codex-rs/app-server/src/code_mode_host.rs create mode 100644 codex-rs/app-server/src/code_mode_host_tests.rs create mode 100644 codex-rs/app-server/tests/suite/v2/code_mode_host.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/transport.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index bf1d6f84ce00..ff64974baa72 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1966,6 +1966,7 @@ dependencies = [ "codex-backend-client", "codex-chatgpt", "codex-cloud-config", + "codex-code-mode", "codex-config", "codex-connectors", "codex-core", @@ -2486,12 +2487,15 @@ name = "codex-code-mode" version = "0.0.0" dependencies = [ "codex-code-mode-protocol", + "codex-http-client", "codex-protocol", + "codex-websocket-client", "deno_core_icudata", "futures", "pretty_assertions", "serde_json", "tokio", + "tokio-tungstenite", "tokio-util", "tracing", "v8", diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index ef18d17b87b0..cf02f023cf69 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -37,6 +37,7 @@ codex-analytics = { workspace = true } codex-agent-extension = { workspace = true } codex-arg0 = { workspace = true } codex-cloud-config = { workspace = true } +codex-code-mode = { workspace = true } codex-config = { workspace = true } codex-connectors = { workspace = true } codex-core = { workspace = true } @@ -105,6 +106,7 @@ tokio = { workspace = true, features = [ tokio-util = { workspace = true } tracing = { workspace = true, features = ["log"] } tracing-subscriber = { workspace = true, features = ["env-filter", "fmt", "json"] } +url = { workspace = true } uuid = { workspace = true, features = ["serde", "v7"] } [target.'cfg(windows)'.dependencies] @@ -138,5 +140,4 @@ shlex = { workspace = true } tar = { workspace = true } tokio-tungstenite = { workspace = true } tracing-opentelemetry = { workspace = true } -url = { workspace = true } wiremock = { workspace = true } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 897631645092..6470c3184140 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -36,6 +36,8 @@ When running with `--listen ws://IP:PORT`, the same listener also serves basic H Websocket transport is currently experimental and unsupported. Do not rely on it for production workloads. +Pass `--code-mode-host wss://HOST/PATH` to connect this app-server process to a remote code-mode host instead of starting a local host. This outbound connection is independent of `--listen` and is shared by the process's threads. Use `ws://` for a local code-mode host. + The unix socket transport is intended for local app-server control-plane clients. `codex app-server proxy` opens exactly one raw stream connection to `$CODEX_HOME/app-server-control/app-server-control.sock` by default, or to `--sock PATH` when provided, and proxies bytes between that socket and stdin/stdout. diff --git a/codex-rs/app-server/src/code_mode_host.rs b/codex-rs/app-server/src/code_mode_host.rs new file mode 100644 index 000000000000..fb3c7e9139f7 --- /dev/null +++ b/codex-rs/app-server/src/code_mode_host.rs @@ -0,0 +1,48 @@ +use clap::Args; +use url::Url; + +/// Selects the code-mode host for a single app-server process. +#[derive(Args, Debug, Clone, Default, PartialEq, Eq)] +pub struct AppServerCodeModeHostArgs { + /// Connect to a remote code-mode host instead of starting a local host. + #[arg( + long = "code-mode-host", + value_name = "WS_URL", + value_parser = parse_websocket_url + )] + pub code_mode_host: Option, +} + +/// Process-scoped transport used to reach the code-mode host. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub enum CodeModeHostTransport { + /// Start and own the default local code-mode host. + #[default] + Local, + /// Share a connection to the specified remote code-mode host. + WebSocket(Url), +} + +impl From for CodeModeHostTransport { + fn from(args: AppServerCodeModeHostArgs) -> Self { + match args.code_mode_host { + Some(url) => Self::WebSocket(url), + None => Self::Local, + } + } +} + +fn parse_websocket_url(value: &str) -> Result { + let url = Url::parse(value).map_err(|error| format!("invalid websocket URL: {error}"))?; + if !matches!(url.scheme(), "ws" | "wss") || url.host_str().is_none() { + return Err("code-mode host URL must use ws:// or wss:// with a host".to_string()); + } + if url.fragment().is_some() { + return Err("code-mode host URL must not contain a fragment".to_string()); + } + Ok(url) +} + +#[cfg(test)] +#[path = "code_mode_host_tests.rs"] +mod tests; diff --git a/codex-rs/app-server/src/code_mode_host_tests.rs b/codex-rs/app-server/src/code_mode_host_tests.rs new file mode 100644 index 000000000000..400406a93ce0 --- /dev/null +++ b/codex-rs/app-server/src/code_mode_host_tests.rs @@ -0,0 +1,51 @@ +use super::AppServerCodeModeHostArgs; +use super::CodeModeHostTransport; +use super::parse_websocket_url; +use pretty_assertions::assert_eq; +use url::Url; + +#[test] +fn websocket_host_accepts_local_and_secure_endpoints() { + for endpoint in ["ws://127.0.0.1:8765", "wss://example.test/code-mode"] { + assert_eq!( + parse_websocket_url(endpoint), + Ok(Url::parse(endpoint).expect("test endpoint should parse")) + ); + } +} + +#[test] +fn websocket_host_rejects_invalid_endpoints() { + for endpoint in [ + "http://127.0.0.1:8765", + "https://example.test/code-mode", + "ws://", + "not a websocket", + "wss://example.test/code-mode#fragment", + ] { + assert!( + parse_websocket_url(endpoint).is_err(), + "invalid code-mode host endpoint should be rejected: {endpoint}" + ); + } +} + +#[test] +fn omitted_websocket_host_selects_local_transport() { + assert_eq!( + CodeModeHostTransport::from(AppServerCodeModeHostArgs::default()), + CodeModeHostTransport::Local + ); +} + +#[test] +fn explicit_websocket_host_selects_remote_transport() { + let url = Url::parse("wss://example.test/code-mode").expect("test endpoint should parse"); + + assert_eq!( + CodeModeHostTransport::from(AppServerCodeModeHostArgs { + code_mode_host: Some(url.clone()), + }), + CodeModeHostTransport::WebSocket(url) + ); +} diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 7be5c055fe8b..f75a1b750c35 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -450,6 +450,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult Self { Self { + code_mode_host_transport: CodeModeHostTransport::Local, plugin_startup_tasks: PluginStartupTasks::Start, remote_control_startup_mode: RemoteControlStartupMode::ResolvePersisted, install_shutdown_signal_handler: true, @@ -534,6 +542,24 @@ pub async fn run_main_with_transport_options( })? } }; + let code_mode_session_provider: Option> = + match &runtime_options.code_mode_host_transport { + CodeModeHostTransport::Local => None, + CodeModeHostTransport::WebSocket(url) => { + if !config.features.enabled(Feature::CodeModeHost) { + return Err(std::io::Error::new( + ErrorKind::InvalidInput, + "remote code-mode host requires the code_mode_host feature to be enabled", + )); + } + Some(Arc::new( + WebSocketCodeModeSessionProvider::with_http_client_factory( + url.to_string(), + config.http_client_factory(), + ), + )) + } + }; let environment_manager = if ignore_user_config { EnvironmentManager::from_env(Some(local_runtime_paths), config.http_client_factory()).await } else { @@ -869,6 +895,7 @@ pub async fn run_main_with_transport_options( session_source, auth_manager, installation_id, + code_mode_session_provider, rpc_transport: analytics_rpc_transport(&transport), remote_control_handle: Some(remote_control_handle.clone()), plugin_startup_tasks: runtime_options.plugin_startup_tasks, diff --git a/codex-rs/app-server/src/main.rs b/codex-rs/app-server/src/main.rs index 7599e498d9d7..4d5ab3f122bf 100644 --- a/codex-rs/app-server/src/main.rs +++ b/codex-rs/app-server/src/main.rs @@ -1,4 +1,5 @@ use clap::Parser; +use codex_app_server::AppServerCodeModeHostArgs; use codex_app_server::AppServerRuntimeOptions; use codex_app_server::AppServerTransport; use codex_app_server::AppServerWebsocketAuthArgs; @@ -22,6 +23,9 @@ struct AppServerArgs { #[command(flatten)] config_overrides: CliConfigOverrides, + #[command(flatten)] + code_mode_host: AppServerCodeModeHostArgs, + /// Transport endpoint URL. Supported values: `stdio://` (default), /// `unix://`, `unix://PATH`, `ws://IP:PORT`, `off`. #[arg( @@ -63,6 +67,7 @@ fn main() -> anyhow::Result<()> { arg0_dispatch_or_else(move |arg0_paths: Arg0DispatchPaths| async move { let AppServerArgs { config_overrides, + code_mode_host, listen, session_source, auth, @@ -80,7 +85,10 @@ fn main() -> anyhow::Result<()> { }; let transport = listen; let auth = auth.try_into_settings()?; - let mut runtime_options = AppServerRuntimeOptions::default(); + let mut runtime_options = AppServerRuntimeOptions { + code_mode_host_transport: code_mode_host.into(), + ..Default::default() + }; #[cfg(debug_assertions)] if disable_plugin_startup_tasks_for_tests { runtime_options.plugin_startup_tasks = PluginStartupTasks::Skip; diff --git a/codex-rs/app-server/src/main_tests.rs b/codex-rs/app-server/src/main_tests.rs index 57d0e5217cce..9eb8d6fd539b 100644 --- a/codex-rs/app-server/src/main_tests.rs +++ b/codex-rs/app-server/src/main_tests.rs @@ -1,7 +1,9 @@ use super::AppServerArgs; use clap::Parser; +use codex_app_server::AppServerTransport; use pretty_assertions::assert_eq; use toml::Value as TomlValue; +use url::Url; #[test] fn app_server_accepts_cli_config_overrides() { @@ -35,3 +37,37 @@ fn app_server_accepts_cli_config_overrides() { ] ); } + +#[test] +fn app_server_accepts_process_scoped_code_mode_host() { + let args = AppServerArgs::try_parse_from([ + "codex-app-server", + "--code-mode-host", + "wss://example.test/code-mode", + "--listen", + "off", + ]) + .expect("parse app-server args"); + + assert_eq!( + args.code_mode_host.code_mode_host, + Some(Url::parse("wss://example.test/code-mode").expect("test endpoint should parse")) + ); + assert_eq!(args.listen, AppServerTransport::Off); + assert_eq!(args.config_overrides.raw_overrides, Vec::::new()); +} + +#[test] +fn app_server_rejects_invalid_code_mode_host() { + for endpoint in [ + "http://127.0.0.1:8765", + "ws://", + "wss://example.test/code-mode#fragment", + ] { + let error = + AppServerArgs::try_parse_from(["codex-app-server", "--code-mode-host", endpoint]) + .expect_err("invalid code-mode host endpoint should fail startup argument parsing"); + + assert_eq!(error.kind(), clap::error::ErrorKind::ValueValidation); + } +} diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index a8935cdd923c..59d7e2b4f7e5 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -63,6 +63,7 @@ use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::experimental_required_message; use codex_arg0::Arg0DispatchPaths; use codex_chatgpt::workspace_settings; +use codex_code_mode::CodeModeSessionProvider; use codex_core::ThreadManager; use codex_core::config::Config; use codex_exec_server::EnvironmentManager; @@ -210,6 +211,7 @@ pub(crate) struct MessageProcessorArgs { pub(crate) session_source: SessionSource, pub(crate) auth_manager: Arc, pub(crate) installation_id: String, + pub(crate) code_mode_session_provider: Option>, pub(crate) rpc_transport: AppServerRpcTransport, pub(crate) remote_control_handle: Option, pub(crate) plugin_startup_tasks: crate::PluginStartupTasks, @@ -233,6 +235,7 @@ impl MessageProcessor { session_source, auth_manager, installation_id, + code_mode_session_provider, rpc_transport, remote_control_handle, plugin_startup_tasks, @@ -253,7 +256,7 @@ impl MessageProcessor { ); let goal_service = Arc::new(GoalService::new()); let thread_manager = Arc::new_cyclic(|thread_manager| { - ThreadManager::new( + let manager = ThreadManager::new( config.as_ref(), auth_manager.clone(), codex_core::build_models_manager(config.as_ref(), auth_manager.clone()), @@ -294,7 +297,11 @@ impl MessageProcessor { outgoing.clone(), thread_state_manager.clone(), )), - ) + ); + match code_mode_session_provider { + Some(provider) => manager.with_code_mode_session_provider(provider), + None => manager, + } }); let models_manager = thread_manager.get_models_manager(); let models_refresh_worker = diff --git a/codex-rs/app-server/src/message_processor_tracing_tests.rs b/codex-rs/app-server/src/message_processor_tracing_tests.rs index e7451cb3fa19..40262c8704dd 100644 --- a/codex-rs/app-server/src/message_processor_tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor_tracing_tests.rs @@ -264,6 +264,7 @@ async fn build_test_processor( session_source: SessionSource::VSCode, auth_manager, installation_id: "11111111-1111-4111-8111-111111111111".to_string(), + code_mode_session_provider: None, rpc_transport: AppServerRpcTransport::Stdio, remote_control_handle: None, plugin_startup_tasks: crate::PluginStartupTasks::Start, diff --git a/codex-rs/app-server/tests/suite/v2/code_mode_host.rs b/codex-rs/app-server/tests/suite/v2/code_mode_host.rs new file mode 100644 index 000000000000..6cfa20492a61 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/code_mode_host.rs @@ -0,0 +1,137 @@ +use std::process::Stdio; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use app_test_support::MockResponsesConfig; +use app_test_support::TestAppServer; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStatus; +use codex_app_server_protocol::UserInput; +use codex_features::Feature; +use core_test_support::responses; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::TempDir; +use tokio::io::AsyncBufReadExt; +use tokio::io::BufReader; +use tokio::process::Command; +use tokio::time::timeout; + +#[cfg(any(target_os = "macos", windows))] +const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 60); +#[cfg(not(any(target_os = "macos", windows)))] +const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 10); + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn app_server_shares_flag_selected_code_mode_host_across_threads() -> Result<()> { + let host_program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let mut websocket_host = Command::new(host_program) + .args(["--listen", "ws://127.0.0.1:0"]) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .kill_on_drop(true) + .spawn() + .context("failed to start websocket code-mode host")?; + let stdout = websocket_host + .stdout + .take() + .context("websocket code-mode host stdout was not captured")?; + let mut lines = BufReader::new(stdout).lines(); + let websocket_url = timeout(DEFAULT_READ_TIMEOUT, lines.next_line()) + .await + .context("timed out waiting for websocket code-mode host URL")?? + .context("websocket code-mode host exited before publishing its URL")?; + + let model_server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_sequence( + &model_server, + vec![ + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_custom_tool_call( + "first-remote-cell", + "exec", + "text('remote app-server host')", + ), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("msg-1", "Done"), + responses::ev_completed("resp-2"), + ]), + responses::sse(vec![ + responses::ev_response_created("resp-3"), + responses::ev_custom_tool_call( + "second-remote-cell", + "exec", + "text('remote app-server host')", + ), + responses::ev_completed("resp-3"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("msg-2", "Done"), + responses::ev_completed("resp-4"), + ]), + ], + ) + .await; + + let codex_home = TempDir::new()?; + MockResponsesConfig::new(&model_server.uri()) + .enable_feature(Feature::CodeModeOnly) + .write(codex_home.path())?; + let original_config = std::fs::read_to_string(codex_home.path().join("config.toml"))?; + let mut app_server = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .with_args(&["--code-mode-host", &websocket_url]) + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) + .await?; + + for prompt in ["run the first remote cell", "run the second remote cell"] { + let thread = app_server + .start_thread(ThreadStartParams::default()) + .await?; + let completed = timeout( + DEFAULT_READ_TIMEOUT, + app_server.start_turn_and_wait_for_completion(TurnStartParams { + thread_id: thread.thread.id, + input: vec![UserInput::Text { + text: prompt.to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }), + ) + .await??; + + assert_eq!(completed.turn.status, TurnStatus::Completed); + } + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 4); + for (request, call_id) in [ + (&requests[1], "first-remote-cell"), + (&requests[3], "second-remote-cell"), + ] { + let output = request.custom_tool_call_output(call_id); + assert_eq!( + output["output"] + .as_array() + .and_then(|items| items.last()) + .cloned(), + Some(json!({ + "type": "input_text", + "text": "remote app-server host", + })) + ); + } + assert_eq!( + std::fs::read_to_string(codex_home.path().join("config.toml"))?, + original_config + ); + + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index ca571ae65c3a..83d5c376716c 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -6,6 +6,7 @@ mod app_read; mod attestation; mod auto_env; mod client_metadata; +mod code_mode_host; mod collaboration_mode_list; #[cfg(unix)] mod command_exec; diff --git a/codex-rs/app-server/tests/suite/v2/remote_control.rs b/codex-rs/app-server/tests/suite/v2/remote_control.rs index b8ec92952a4d..c8805062b890 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_control.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_control.rs @@ -233,6 +233,7 @@ async fn explicit_remote_control_startup_fails_when_disabled_by_requirements() - plugin_startup_tasks: PluginStartupTasks::Skip, remote_control_startup_mode: RemoteControlStartupMode::EnabledEphemeral, install_shutdown_signal_handler: false, + ..Default::default() }, ), ) diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index c01a94b2d377..034c27fdda38 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -516,6 +516,9 @@ struct AppServerCommand { #[command(subcommand)] subcommand: Option, + #[command(flatten)] + code_mode_host: codex_app_server::AppServerCodeModeHostArgs, + /// Error out when config.toml contains fields that are not recognized by this version of Codex. #[arg(long = "strict-config", default_value_t = false)] strict_config: bool, @@ -1099,6 +1102,7 @@ async fn cli_main( Some(Subcommand::AppServer(app_server_cli)) => { let AppServerCommand { subcommand, + code_mode_host, strict_config: app_server_strict_config, listen, stdio, @@ -1122,6 +1126,7 @@ async fn cli_main( }; let auth = auth.try_into_settings()?; let runtime_options = codex_app_server::AppServerRuntimeOptions { + code_mode_host_transport: code_mode_host.into(), remote_control_startup_mode: match (remote_control, remote_control_disabled) { (true, _) => { @@ -3817,6 +3822,50 @@ mod tests { assert!(err.to_string().contains("is empty")); } + #[test] + fn app_server_code_mode_host_url_parses_independently_of_listen_transport() { + let app_server = app_server_from_args( + [ + "codex", + "app-server", + "--code-mode-host", + "wss://example.test/code-mode", + "--listen", + "ws://127.0.0.1:4500", + ] + .as_ref(), + ); + + assert_eq!( + app_server.code_mode_host.code_mode_host, + Some( + url::Url::parse("wss://example.test/code-mode") + .expect("test endpoint should parse") + ) + ); + assert_eq!( + app_server.listen, + codex_app_server::AppServerTransport::WebSocket { + bind_address: "127.0.0.1:4500".parse().expect("valid socket address"), + } + ); + } + + #[test] + fn app_server_rejects_invalid_code_mode_host_urls() { + for endpoint in [ + "http://127.0.0.1:8765", + "ws://", + "wss://example.test/code-mode#fragment", + ] { + let error = + MultitoolCli::try_parse_from(["codex", "app-server", "--code-mode-host", endpoint]) + .expect_err("invalid code-mode host endpoint should fail argument parsing"); + + assert_eq!(error.kind(), clap::error::ErrorKind::ValueValidation); + } + } + #[test] fn app_server_listen_websocket_url_parses() { let app_server = app_server_from_args( diff --git a/codex-rs/code-mode-host/src/peer.rs b/codex-rs/code-mode-host/src/peer.rs index 14ff068e7b8d..79120518025f 100644 --- a/codex-rs/code-mode-host/src/peer.rs +++ b/codex-rs/code-mode-host/src/peer.rs @@ -13,6 +13,7 @@ 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::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::WireResult; @@ -25,7 +26,6 @@ use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; const CELL_MESSAGE_CAPACITY: usize = 128; -const MAX_PENDING_DELEGATE_CALLS: usize = 256; pub(super) struct HostPeer { outgoing_tx: mpsc::Sender, diff --git a/codex-rs/code-mode-protocol/src/host/mod.rs b/codex-rs/code-mode-protocol/src/host/mod.rs index e448b670ea13..5c81b1d4a028 100644 --- a/codex-rs/code-mode-protocol/src/host/mod.rs +++ b/codex-rs/code-mode-protocol/src/host/mod.rs @@ -11,6 +11,9 @@ mod message; mod payload; mod types; +/// Maximum number of unresolved delegate callbacks allowed per host connection. +pub const MAX_PENDING_DELEGATE_CALLS: usize = 1_024; + pub use codec::EncodedFrame; pub use codec::FramedReader; pub use codec::FramedWriter; diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index be1442418137..56c77917628b 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/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 } diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 6d332bf0c440..c4266a4a55fb 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -10,6 +10,7 @@ pub(crate) type TaskFailureHandler = std::sync::Arc, } +/// Creates code-mode sessions backed by one shared remote WebSocket connection. +pub struct WebSocketCodeModeSessionProvider { + host: Arc, +} + enum ProviderState { - OwnedProcess(Arc), + OwnedProcess(Arc), 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> { + fn process_host(&self) -> Option> { 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 = 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, + ) -> CodeModeSessionProviderFuture<'a> { + Box::pin(create_host_session(delegate, Arc::clone(&self.host))) + } +} + +async fn create_host_session( + delegate: Arc, + host: Arc, +) -> Result, 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>>, - 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), + 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), - spawn_permit: Semaphore::new(/*permits*/ 1), + 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, + host: Arc, delegate: Arc, state: StdMutex, next_generation: AtomicU64, @@ -186,26 +257,23 @@ struct SessionInner { retired_cleanups: StdMutex>, } -/// 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, } 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, - process_host: Arc, - ) -> Self { + fn with_host(delegate: Arc, host: Arc) -> 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>>, ) { - 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)) diff --git a/codex-rs/code-mode/src/remote_session/connection.rs b/codex-rs/code-mode/src/remote_session/connection.rs index 18ba5e09743c..ab386e42d1cf 100644 --- a/codex-rs/code-mode/src/remote_session/connection.rs +++ b/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::(); pub(super) enum ConnectionError { Spawn { @@ -96,7 +106,7 @@ struct CallerCancellation { } struct ConnectionSupervisor { - child: Child, + owner: ConnectionOwner, event_tx: mpsc::Sender, cancellation: CancellationToken, alive: Arc, @@ -106,6 +116,11 @@ struct ConnectionSupervisor { writer_task: JoinHandle>, } +enum ConnectionOwner { + Process(Box), + 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 { + 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 { 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::() + .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 => {} } } } diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs b/codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs index c79270661d08..94be7b1b645c 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs +++ b/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)>, pub(super) closed_cells: Vec, @@ -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( diff --git a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs index 97ed9b553e32..9f0f3f762500 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs +++ b/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::(&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(); diff --git a/codex-rs/code-mode/src/remote_session/connection/reader.rs b/codex-rs/code-mode/src/remote_session/connection/reader.rs index fbfc94de4278..8fb7cd8b294d 100644 --- a/codex-rs/code-mode/src/remote_session/connection/reader.rs +++ b/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, + mut reader: ConnectionReader, events: mpsc::Sender, cancellation: CancellationToken, ) -> Result<(), String> { loop { let message = tokio::select! { _ = cancellation.cancelled() => return Ok(()), - result = reader.read::() => result, + result = reader.read() => result, }; let message = match message { Ok(Some(message)) => message, diff --git a/codex-rs/code-mode/src/remote_session/connection/transport.rs b/codex-rs/code-mode/src/remote_session/connection/transport.rs new file mode 100644 index 000000000000..9ce520b5f2e4 --- /dev/null +++ b/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), + WebSocket(SplitStream), +} + +pub(super) enum ConnectionWriter { + Stdio(FramedWriter), + WebSocket(SplitSink), +} + +impl ConnectionReader { + pub(super) async fn read(&mut self) -> io::Result> { + 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}" + )) + }) + } + } + } +} diff --git a/codex-rs/code-mode/src/remote_session_tests.rs b/codex-rs/code-mode/src/remote_session_tests.rs index e48ac2d71ef0..68f03a3a2cb4 100644 --- a/codex-rs/code-mode/src/remote_session_tests.rs +++ b/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::(&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(); diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index a4f5bebc7d4c..105505c1a4e5 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -396,6 +396,18 @@ impl ThreadManager { } } + /// Replaces the process-wide provider before this manager is shared with threads. + pub fn with_code_mode_session_provider( + mut self, + provider: Arc, + ) -> Self { + let Some(state) = Arc::get_mut(&mut self.state) else { + unreachable!("code-mode session provider must be set before thread manager is shared"); + }; + state.code_mode_session_provider = provider; + self + } + pub(crate) fn with_code_mode_host_program_for_tests(mut self, host_program: PathBuf) -> Self { let Some(state) = Arc::get_mut(&mut self.state) else { unreachable!("new thread manager state should not be shared"); diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index c08a7ab515bd..0d29bcbe085f 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -428,12 +428,14 @@ async fn code_mode_session_provider_is_shared_across_threads() { config.cwd = config.codex_home.abs(); std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + let provider: Arc = Arc::new(InProcessCodeModeSessionProvider); let manager = ThreadManager::with_models_provider_and_home_for_tests( CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .with_code_mode_session_provider(Arc::clone(&provider)); let first = manager .start_thread(StartThreadOptions::new(config.clone())) .await @@ -456,6 +458,7 @@ async fn code_mode_session_provider_is_shared_across_threads() { .code_mode_service .session_provider(); assert!(Arc::ptr_eq(&first_provider, &second_provider)); + assert!(Arc::ptr_eq(&first_provider, &provider)); assert!(Arc::ptr_eq( &first_provider, &manager.state.code_mode_session_provider