openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 17:28:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
## 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
This commit is contained in:
Channing Conger
2026-07-24 04:37:01 +00:00
committed by copyberry
parent 0dfa778dae
commit f61b51ddd9
28 changed files with 977 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+4
View File
@@ -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",
codex-rs/app-server/Cargo.toml
+2 -1
View File
@@ -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 }
codex-rs/app-server/README.md
+2
View File
@@ -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.
codex-rs/app-server/src/code_mode_host.rs
+48
View File
@@ -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<Url>,
}
/// 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<AppServerCodeModeHostArgs> 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<Url, String> {
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;
codex-rs/app-server/src/code_mode_host_tests.rs
+51
View File
@@ -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)
);
}
codex-rs/app-server/src/in_process.rs
+1
View File
@@ -450,6 +450,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
session_source: args.session_source,
auth_manager,
installation_id,
code_mode_session_provider: None,
rpc_transport: AppServerRpcTransport::InProcess,
remote_control_handle: None,
plugin_startup_tasks: crate::PluginStartupTasks::Start,
codex-rs/app-server/src/lib.rs
+28 -1
View File
@@ -2,6 +2,8 @@
#![deny(clippy::print_stdout, clippy::print_stderr)]
use codex_arg0::Arg0DispatchPaths;
use codex_code_mode::CodeModeSessionProvider;
use codex_code_mode::WebSocketCodeModeSessionProvider;
use codex_config::ConfigLayerStackOrdering;
use codex_config::LoaderOverrides;
use codex_config::NoopThreadConfigLoader;
@@ -60,6 +62,7 @@ use codex_core::check_execpolicy_for_warnings;
use codex_core::config::find_codex_home;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::ExecServerRuntimePaths;
use codex_features::Feature;
use codex_feedback::CodexFeedback;
use codex_protocol::protocol::SessionSource;
use codex_rollout::state_db as rollout_state_db;
@@ -85,6 +88,7 @@ mod app_server_tracing;
mod attestation;
mod auth_mode;
mod bespoke_event_handling;
mod code_mode_host;
mod command_exec;
mod config_layer;
mod config_manager;
@@ -116,6 +120,8 @@ mod thread_state;
mod thread_status;
mod transport;
pub use crate::code_mode_host::AppServerCodeModeHostArgs;
pub use crate::code_mode_host::CodeModeHostTransport;
pub use crate::error_code::INPUT_TOO_LARGE_ERROR_CODE;
pub use crate::error_code::INVALID_PARAMS_ERROR_CODE;
pub use crate::transport::AppServerTransport;
@@ -428,8 +434,9 @@ pub enum PluginStartupTasks {
Skip,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AppServerRuntimeOptions {
pub code_mode_host_transport: CodeModeHostTransport,
pub plugin_startup_tasks: PluginStartupTasks,
pub remote_control_startup_mode: RemoteControlStartupMode,
pub install_shutdown_signal_handler: bool,
@@ -438,6 +445,7 @@ pub struct AppServerRuntimeOptions {
impl Default for AppServerRuntimeOptions {
fn default() -> 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<Arc<dyn CodeModeSessionProvider>> =
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,
codex-rs/app-server/src/main.rs
+9 -1
View File
@@ -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;
codex-rs/app-server/src/main_tests.rs
+36
View File
@@ -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::<String>::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);
}
}
codex-rs/app-server/src/message_processor.rs
+9 -2
View File
@@ -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<AuthManager>,
pub(crate) installation_id: String,
pub(crate) code_mode_session_provider: Option<Arc<dyn CodeModeSessionProvider>>,
pub(crate) rpc_transport: AppServerRpcTransport,
pub(crate) remote_control_handle: Option<RemoteControlHandle>,
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 =
codex-rs/app-server/src/message_processor_tracing_tests.rs
+1
View File
@@ -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,
codex-rs/app-server/tests/suite/v2/code_mode_host.rs
+137
View File
@@ -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(())
}
codex-rs/app-server/tests/suite/v2/mod.rs
+1
View File
@@ -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;
codex-rs/app-server/tests/suite/v2/remote_control.rs
+1
View File
@@ -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()
},
),
)
codex-rs/cli/src/main.rs
+49
View File
@@ -516,6 +516,9 @@ struct AppServerCommand {
#[command(subcommand)]
subcommand: Option<AppServerSubcommand>,
#[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(
codex-rs/code-mode-host/src/peer.rs
+1 -1
View File
@@ -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<EncodedFrame>,
codex-rs/code-mode-protocol/src/host/mod.rs
+3
View File
@@ -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;
codex-rs/code-mode/Cargo.toml
+4 -1
View File
@@ -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
+1
View File
@@ -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
+95 -27
View File
@@ -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
+112 -18
View File
@@ -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
+24 -7
View File
@@ -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
+70
View File
@@ -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
+3 -5
View File
@@ -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
+103
View File
@@ -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
+166
View File
@@ -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();
codex-rs/core/src/thread_manager.rs
+12
View File
@@ -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<dyn CodeModeSessionProvider>,
) -> 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");
codex-rs/core/src/thread_manager_tests.rs
+4 -1
View File
@@ -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<dyn CodeModeSessionProvider> = 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