diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 12036f3a3b..194dbbfa0d 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2760,8 +2760,10 @@ dependencies = [ "codex-http-client", "codex-otel", "codex-otel-trace-websocket", + "codex-plugin-host", "codex-protocol", "codex-utils-cargo-bin", + "codex-utils-process", "futures", "pretty_assertions", "prost", @@ -2800,6 +2802,7 @@ version = "0.0.0" dependencies = [ "base64 0.22.1", "codex-code-mode-protocol", + "codex-plugin-host", "codex-protocol", "deno_core_icudata", "futures", @@ -4392,6 +4395,7 @@ dependencies = [ "futures", "libc", "pretty_assertions", + "serde", "serde_json", "tempfile", "tokio", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 830308f07e..5583fde4ff 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -209,6 +209,7 @@ codex-cloud-tasks-mock-client = { path = "cloud-tasks-mock-client" } codex-code-mode = { path = "code-mode" } codex-code-mode-protocol = { path = "code-mode-protocol" } codex-code-mode-runtime = { path = "code-mode-runtime" } +codex-plugin-host = { path = "plugin-host" } codex-home = { path = "codex-home" } codex-http-client = { path = "http-client" } codex-websocket-auth = { path = "websocket-auth" } diff --git a/codex-rs/code-mode-host/AGENTS.md b/codex-rs/code-mode-host/AGENTS.md new file mode 100644 index 0000000000..3efe7dc2dc --- /dev/null +++ b/codex-rs/code-mode-host/AGENTS.md @@ -0,0 +1,3 @@ +# Tests + +Do not use timeouts or sleeps to synchronize tests. Wait for explicit signals or observable state; let the test runner enforce test-case timeouts. diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index 4b0d51c1a9..98e45c5e0b 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -38,7 +38,9 @@ tracing-subscriber = { workspace = true } uuid = { workspace = true, features = ["v4"] } [dev-dependencies] +codex-plugin-host = { workspace = true } codex-code-mode = { workspace = true } codex-utils-cargo-bin = { workspace = true } +codex-utils-process = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index cc797ffa47..fc7cedb848 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -18,6 +18,13 @@ const OTEL_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5); #[derive(Debug, Parser)] struct Cli { + #[arg( + long, + hide = true, + conflicts_with_all = ["listen", "otel_trace_listen", "otel_trace_exporter"] + )] + codex_ui_plugin_process: bool, + /// Transport endpoint: `stdio`, `stdio://`, `grpc+stdio://`, or `grpc://IP:PORT`. #[arg( long, @@ -37,8 +44,7 @@ struct Cli { } #[tokio::main(flavor = "current_thread")] -async fn main() -> anyhow::Result<()> { - let cli = Cli::parse(); +async fn run(cli: Cli) -> anyhow::Result<()> { let mut trace_transport = if let Some(trace_listen) = cli.otel_trace_listen.as_deref() { Some(TraceWebSocket::start(trace_listen).await?) } else { @@ -88,6 +94,14 @@ async fn main() -> anyhow::Result<()> { result } +fn main() -> anyhow::Result<()> { + let cli = Cli::parse(); + if cli.codex_ui_plugin_process { + return codex_code_mode_runtime::run_ui_sandbox().map_err(Into::into); + } + run(cli) +} + fn build_trace_provider(endpoint: &str) -> anyhow::Result { OtelProvider::try_new(&OtelSettings { http_client_factory: codex_http_client::HttpClientFactory::new( diff --git a/codex-rs/code-mode-host/tests/ui_plugin_wire.rs b/codex-rs/code-mode-host/tests/ui_plugin_wire.rs new file mode 100644 index 0000000000..2076a583d0 --- /dev/null +++ b/codex-rs/code-mode-host/tests/ui_plugin_wire.rs @@ -0,0 +1,506 @@ +//! Exercises the child protocol without the parent supervisor or session APIs. + +use anyhow::Context; +use codex_plugin_host::PluginProcess; +use codex_plugin_host::protocol::MAX_PENDING_PROMISES; +use codex_plugin_host::protocol::MAX_SCRIPT_ERROR_BYTES; +use codex_plugin_host::protocol::Packet; +use futures::SinkExt; +use futures::StreamExt; +use pretty_assertions::assert_eq; +use serde_json::json; + +#[tokio::test] +async fn child_initializes_invokes_and_closes_an_isolate() -> anyhow::Result<()> { + let program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let mut child = + PluginProcess::spawn(&program, &tokio_util::sync::CancellationToken::new()).await?; + let result = async { + let source = format!("/*{}*/ export default () => 'x'.repeat(5000);", "p".repeat(40 * 1024)); + for (message, expected) in [ + (json!({"type": "initialize", "source": source}), json!({"type": "initialized"})), + (json!({"type": "dispatch", "id": 0, "export": "default", "returnMode": "maybeAsync"}), json!({"type": "reply", "id": 0, "value": "x".repeat(5000)})), + (json!({"type": "close"}), json!({"type": "closed"})), + ] { + send(&mut child, /*instance*/ 7, message).await?; + assert_eq!(receive(&mut child).await?, (7, expected)); + } + anyhow::Ok(()) + } + .await; + child.close().await?; + result +} + +async fn send( + child: &mut PluginProcess, + instance: u64, + message: serde_json::Value, +) -> anyhow::Result<()> { + let (input, _) = child.pipes(); + input + .send(serde_json::to_vec(&Packet { instance, message })?.into()) + .await?; + Ok(()) +} + +async fn receive(child: &mut PluginProcess) -> anyhow::Result<(u64, serde_json::Value)> { + let (_, output) = child.pipes(); + let packet: Packet = serde_json::from_slice( + &output + .next() + .await + .context("plugin child closed stdout")??, + )?; + Ok((packet.instance, packet.message)) +} + +#[tokio::test] +async fn promise_errors_close_the_isolate_in_both_return_modes() -> anyhow::Result<()> { + let program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let mut child = + PluginProcess::spawn(&program, &tokio_util::sync::CancellationToken::new()).await?; + let result = async { + let source = r#" + export function run() { + const promise = Promise.reject(Error('tracked failure')); + promise.catch(() => {}); + return promise; + } + "#; + for (instance, return_mode, code) in [ + (1, "sync", "invalid_result"), + (2, "maybeAsync", "execution_failed"), + ] { + send( + &mut child, + instance, + json!({"type": "initialize", "source": source}), + ) + .await?; + assert_eq!( + receive(&mut child).await?, + (instance, json!({"type": "initialized"})) + ); + send( + &mut child, + instance, + json!({"type": "dispatch", "id": 0, "export": "run", "returnMode": return_mode}), + ) + .await?; + let (id, error) = receive(&mut child).await?; + assert_eq!( + (id, error["type"].as_str(), error["code"].as_str()), + (instance, Some("error"), Some(code)) + ); + if return_mode == "maybeAsync" { + assert!( + error["scriptError"]["message"] + .as_str() + .is_some_and(|diagnostic| diagnostic.contains("tracked failure")) + ); + } + assert_eq!( + receive(&mut child).await?, + (instance, json!({"type": "closed"})) + ); + } + anyhow::Ok(()) + } + .await; + child.close().await?; + result +} + +#[tokio::test] +async fn cancellation_closes_an_isolate_with_pending_work() -> anyhow::Result<()> { + let program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let mut child = + PluginProcess::spawn(&program, &tokio_util::sync::CancellationToken::new()).await?; + let result = async { + let source = "export function run() { return new Promise(() => {}); }"; + send( + &mut child, + /*instance*/ 1, + json!({"type": "initialize", "source": source}), + ) + .await?; + assert_eq!( + receive(&mut child).await?, + (1, json!({"type": "initialized"})) + ); + send( + &mut child, + /*instance*/ 1, + json!({"type": "dispatch", "id": 0, "export": "run", "returnMode": "maybeAsync"}), + ) + .await?; + assert_eq!( + receive(&mut child).await?, + (1, json!({"type": "pending", "id": 0})) + ); + send(&mut child, /*instance*/ 1, json!({"type": "close"})).await?; + assert_eq!(receive(&mut child).await?, (1, json!({"type": "closed"}))); + anyhow::Ok(()) + } + .await; + child.close().await?; + result +} + +#[tokio::test] +async fn sync_dispatch_settles_pending_invocations() -> anyhow::Result<()> { + let program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let mut child = + PluginProcess::spawn(&program, &tokio_util::sync::CancellationToken::new()).await?; + let result = async { + let source = r#" + const resolvers = []; + export function wait() { return new Promise(resolve => resolvers.push(resolve)); } + export function finish() { + for (const resolve of resolvers) resolve('done'); + return 'finished'; + } + "#; + send( + &mut child, + /*instance*/ 1, + json!({"type": "initialize", "source": source}), + ) + .await?; + assert_eq!( + receive(&mut child).await?, + (1, json!({"type": "initialized"})) + ); + let sync_id = MAX_PENDING_PROMISES as u64; + for id in 0..sync_id { + send( + &mut child, + /*instance*/ 1, + json!({"type": "dispatch", "id": id, "export": "wait", "returnMode": "maybeAsync"}), + ) + .await?; + assert_eq!( + receive(&mut child).await?, + (1, json!({"type": "pending", "id": id})) + ); + } + send( + &mut child, + /*instance*/ 1, + json!({"type": "dispatch", "id": sync_id, "export": "finish", "returnMode": "sync"}), + ) + .await?; + assert_eq!( + receive(&mut child).await?, + ( + 1, + json!({"type": "reply", "id": sync_id, "value": "finished"}) + ) + ); + let mut resolved = Vec::new(); + for _ in 0..sync_id { + let (instance, reply) = receive(&mut child).await?; + let id = reply["id"].as_u64().unwrap(); + assert_eq!( + (instance, reply), + (1, json!({"type": "reply", "id": id, "value": "done"})) + ); + resolved.push(id); + } + resolved.sort_unstable(); + assert_eq!(resolved, (0..sync_id).collect::>()); + anyhow::Ok(()) + } + .await; + child.close().await?; + result +} + +#[tokio::test] +async fn syntax_and_execution_failures_preserve_bounded_script_errors() -> anyhow::Result<()> { + let program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let mut child = + PluginProcess::spawn(&program, &tokio_util::sync::CancellationToken::new()).await?; + let result = async { + send( + &mut child, + /*instance*/ 1, + json!({"type": "initialize", "source": "export default => ;"}), + ) + .await?; + let (id, error) = receive(&mut child).await?; + assert_eq!(id, 1); + assert_eq!(error["code"], "invalid_program"); + assert_eq!(error["scriptError"]["line"], 1); + assert!(error["scriptError"]["column"].is_number()); + assert_eq!(error["scriptError"]["truncated"], false); + assert_eq!( + receive(&mut child).await?, + (1, json!({"type": "closed"})) + ); + + let source = format!( + r#"export function fail() {{ throw Error('xxx\u001b[31m' + 'x'.repeat({MAX_SCRIPT_ERROR_BYTES})); }}"# + ); + send( + &mut child, + /*instance*/ 2, + json!({"type": "initialize", "source": source}), + ) + .await?; + assert_eq!( + receive(&mut child).await?, + (2, json!({"type": "initialized"})) + ); + send( + &mut child, + /*instance*/ 2, + json!({"type": "dispatch", "id": 0, "export": "fail", "returnMode": "sync"}), + ) + .await?; + let (id, error) = receive(&mut child).await?; + assert_eq!(id, 2); + let message = error["scriptError"]["message"].as_str().unwrap(); + assert!(message.contains("Error: xxx\u{001b}[31m")); + assert!(message.len() <= MAX_SCRIPT_ERROR_BYTES); + assert_eq!(error["scriptError"]["truncated"], true); + assert_eq!( + receive(&mut child).await?, + (2, json!({"type": "closed"})) + ); + anyhow::Ok(()) + } + .await; + child.close().await?; + result +} + +#[tokio::test] +async fn loading_restrictions_close_the_isolate() -> anyhow::Result<()> { + let program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let mut child = + PluginProcess::spawn(&program, &tokio_util::sync::CancellationToken::new()).await?; + let result = async { + let cases = [ + ( + "import 'forbidden';", + "invalid_program", + "module imports are unavailable", + ), + ( + "await import('forbidden');", + "execution_failed", + "Not supported", + ), + ("eval('1 + 1');", "execution_failed", "EvalError"), + ( + "new Function('return 1')();", + "execution_failed", + "EvalError", + ), + ( + "WebAssembly;", + "execution_failed", + "WebAssembly is not defined", + ), + ( + "new ArrayBuffer(1);", + "execution_failed", + "ArrayBuffer is not defined", + ), + ( + "new Uint8Array(1);", + "execution_failed", + "Uint8Array is not defined", + ), + ]; + for (instance, (source, code, diagnostic)) in (1_u64..).zip(cases) { + send( + &mut child, + instance, + json!({"type": "initialize", "source": source}), + ) + .await?; + let (id, error) = receive(&mut child) + .await + .with_context(|| format!("loading {source}"))?; + assert_eq!( + (id, error["type"].as_str(), error["code"].as_str()), + (instance, Some("error"), Some(code)), + "{source}" + ); + assert!( + error["scriptError"]["message"] + .as_str() + .is_some_and(|message| message.contains(diagnostic)), + "{source}: {error}" + ); + assert_eq!( + receive(&mut child).await?, + (instance, json!({"type": "closed"})), + "{source}" + ); + } + anyhow::Ok(()) + } + .await; + let exit = child.close().await?; + result.with_context(|| { + format!( + "plugin helper exited with {}; stderr: {:?}", + exit.status, + String::from_utf8_lossy(&exit.stderr_tail) + ) + }) +} + +#[tokio::test] +async fn isolate_failures_do_not_affect_healthy_peers() -> anyhow::Result<()> { + const HEALTHY: u64 = 1; + const REJECTING: u64 = 2; + const LOOPING: u64 = 3; + + let program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let mut child = + PluginProcess::spawn(&program, &tokio_util::sync::CancellationToken::new()).await?; + let result = async { + let source = r#" + let n = 0; + export function count() { return String(++n); } + export function caught() { + const p = Promise.reject('handled'); + p.catch(() => {}); + return 'ok'; + } + export function detached() { + Promise.resolve().then(() => { throw Error('detached failure'); }); + return 'ok'; + } + "#; + for instance in [HEALTHY, REJECTING] { + send( + &mut child, + instance, + json!({"type": "initialize", "source": source}), + ) + .await?; + assert_eq!( + receive(&mut child).await?, + (instance, json!({"type": "initialized"})) + ); + } + for (instance, export, expected) in [ + (HEALTHY, "count", "1"), + (HEALTHY, "count", "2"), + (REJECTING, "caught", "ok"), + (REJECTING, "count", "1"), + ] { + send( + &mut child, + instance, + json!({"type": "dispatch", "id": 0, "export": export, "returnMode": "maybeAsync"}), + ) + .await?; + assert_eq!( + receive(&mut child).await?, + ( + instance, + json!({"type": "reply", "id": 0, "value": expected}) + ) + ); + } + + send( + &mut child, + REJECTING, + json!({"type": "dispatch", "id": 0, "export": "detached", "returnMode": "maybeAsync"}), + ) + .await?; + let (id, error) = receive(&mut child).await?; + assert_eq!(id, REJECTING); + assert_eq!(error["code"], "execution_failed"); + assert!( + error["scriptError"]["message"] + .as_str() + .unwrap() + .contains("detached failure") + ); + assert_eq!( + receive(&mut child).await?, + (REJECTING, json!({"type": "closed"})) + ); + + let source = "export function loop() { while (true) {} }"; + send( + &mut child, + LOOPING, + json!({"type": "initialize", "source": source}), + ) + .await?; + assert_eq!( + receive(&mut child).await?, + (LOOPING, json!({"type": "initialized"})) + ); + send( + &mut child, + LOOPING, + json!({"type": "dispatch", "id": 0, "export": "loop", "returnMode": "maybeAsync"}), + ) + .await?; + send( + &mut child, + HEALTHY, + json!({"type": "dispatch", "id": 0, "export": "count", "returnMode": "maybeAsync"}), + ) + .await?; + assert_eq!( + receive(&mut child).await?, + ( + LOOPING, + json!({"type": "error", "code": "execution_timed_out", "scriptError": null}) + ) + ); + assert_eq!( + receive(&mut child).await?, + (LOOPING, json!({"type": "closed"})) + ); + assert_eq!( + receive(&mut child).await?, + (HEALTHY, json!({"type": "reply", "id": 0, "value": "3"})) + ); + anyhow::Ok(()) + } + .await; + child.close().await?; + result +} + +#[tokio::test] +async fn closed_instance_can_be_replaced_at_capacity() -> anyhow::Result<()> { + let program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let mut child = + PluginProcess::spawn(&program, &tokio_util::sync::CancellationToken::new()).await?; + let result = async { + let limit = codex_plugin_host::MAX_INSTANCES as u64; + let initialize = json!({"type":"initialize", "source":"export default () => 'ok'"}); + for instance in 0..limit { + send(&mut child, instance, initialize.clone()).await?; + assert_eq!( + receive(&mut child).await?, + (instance, json!({"type":"initialized"})) + ); + } + send(&mut child, /*instance*/ 0, json!({"type":"close"})).await?; + assert_eq!(receive(&mut child).await?, (0, json!({"type":"closed"}))); + send(&mut child, limit, initialize).await?; + assert_eq!( + receive(&mut child).await?, + (limit, json!({"type":"initialized"})) + ); + anyhow::Ok(()) + } + .await; + child.close().await?; + result +} diff --git a/codex-rs/code-mode-host/tests/ui_sandbox.rs b/codex-rs/code-mode-host/tests/ui_sandbox.rs new file mode 100644 index 0000000000..61fee518fb --- /dev/null +++ b/codex-rs/code-mode-host/tests/ui_sandbox.rs @@ -0,0 +1,28 @@ +#[test] +fn plugin_arguments_are_parsed_before_startup() -> Result<(), Box> { + let program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + for arguments in [ + ["--help", "--codex-ui-plugin-process"], + ["--codex-ui-plugin-process", "--help"], + ] { + let output = codex_utils_process::background_command(&program) + .args(arguments) + .output()?; + assert!(output.status.success()); + assert!(String::from_utf8(output.stdout)?.contains("Usage:")); + } + for arguments in [ + ["--unknown", "--codex-ui-plugin-process"], + ["--codex-ui-plugin-process", "--unknown"], + ["--listen=stdio", "--codex-ui-plugin-process"], + ["--codex-ui-plugin-process", "--listen=stdio"], + ] { + let output = codex_utils_process::background_command(&program) + .args(arguments) + .output()?; + assert!(!output.status.success()); + assert!(output.stdout.is_empty()); + assert!(String::from_utf8(output.stderr)?.contains("error:")); + } + Ok(()) +} diff --git a/codex-rs/code-mode-runtime/Cargo.toml b/codex-rs/code-mode-runtime/Cargo.toml index 5a4d9e1d09..786fad083d 100644 --- a/codex-rs/code-mode-runtime/Cargo.toml +++ b/codex-rs/code-mode-runtime/Cargo.toml @@ -15,13 +15,14 @@ workspace = true [dependencies] base64 = { workspace = true } codex-code-mode-protocol = { workspace = true } +codex-plugin-host = { workspace = true } codex-protocol = { workspace = true } deno_core_icudata = { workspace = true } -futures = { workspace = true } +futures = { workspace = true, features = ["executor", "std"] } opentelemetry = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } -tokio-util = { workspace = true, features = ["rt"] } +tokio-util = { workspace = true, features = ["codec", "compat", "rt"] } tracing = { workspace = true } v8 = { workspace = true, features = ["v8_enable_sandbox"] } diff --git a/codex-rs/code-mode-runtime/src/lib.rs b/codex-rs/code-mode-runtime/src/lib.rs index 819636ef0d..0a3e29775e 100644 --- a/codex-rs/code-mode-runtime/src/lib.rs +++ b/codex-rs/code-mode-runtime/src/lib.rs @@ -2,11 +2,13 @@ mod cell_actor; mod runtime; mod service; mod session_runtime; +mod ui; mod v8_init; pub(crate) type TaskFailureHandler = std::sync::Arc; pub use codex_code_mode_protocol::*; pub use service::InProcessCodeModeSession; +pub use ui::run as run_ui_sandbox; pub use v8_init::V8JitMode; pub use v8_init::initialize_v8; diff --git a/codex-rs/code-mode-runtime/src/ui/engine.rs b/codex-rs/code-mode-runtime/src/ui/engine.rs new file mode 100644 index 0000000000..ea8d2715fd --- /dev/null +++ b/codex-rs/code-mode-runtime/src/ui/engine.rs @@ -0,0 +1,278 @@ +//! Persistent module state with explicit microtask pumping and no dynamic executable loading. + +use std::collections::HashMap; +use std::io; + +use codex_plugin_host::protocol::*; + +use super::failures; +use super::transport::Transport; + +pub(super) struct Engine { + context: v8::Global, + namespace: Option>, + pending_promises: HashMap>, + pub(super) transport: Transport, + entered: bool, + isolate: v8::OwnedIsolate, +} + +impl Engine { + /// Construct an isolate; user code runs only in watchdog-protected turns. + pub(super) fn new(transport: Transport) -> io::Result { + // V8 heap exhaustion is fatal to the shared helper; the parent survives and can restart it. + let mut isolate = + v8::Isolate::new(v8::CreateParams::default().heap_limits(0, 64 * 1024 * 1024)); + isolate.set_microtasks_policy(v8::MicrotasksPolicy::Explicit); + failures::install(&mut isolate); + // V8 rejects dynamic imports when no host loader is installed. + isolate.set_allow_wasm_code_generation_callback(reject_wasm); + let context = { + v8::scope!(let scope, &mut isolate); + let context = v8::Context::new(scope, Default::default()); + context.set_allow_generation_from_strings(false); + let scope = &mut v8::ContextScope::new(scope, context); + let global = context.global(scope); + // Binary buffers are omitted until the sandbox-enabled allocator supports quotas. + for name in [ + "console", + "ArrayBuffer", + "SharedArrayBuffer", + "Atomics", + "WebAssembly", + "DataView", + "Int8Array", + "Uint8Array", + "Uint8ClampedArray", + "Int16Array", + "Uint16Array", + "Int32Array", + "Uint32Array", + "Float16Array", + "Float32Array", + "Float64Array", + "BigInt64Array", + "BigUint64Array", + ] { + let key = v8::String::new(scope, name) + .ok_or_else(|| io::Error::other("V8 allocation failed"))?; + if global.delete(scope, key.into()) != Some(true) { + return Err(io::Error::other("failed to restrict UI globals")); + } + } + v8::Global::new(scope, context) + }; + let mut runtime = Self { + entered: true, + transport, + isolate, + context, + namespace: None, + pending_promises: HashMap::new(), + }; + // SAFETY: no handle scopes remain; the executor enters only the isolate being polled. + unsafe { + runtime.isolate.exit(); + } + runtime.entered = false; + Ok(runtime) + } + + pub(super) fn handle(&self) -> v8::IsolateHandle { + self.isolate.thread_safe_handle() + } + + pub(super) fn turn( + &mut self, + message: HostMessage, + ) -> Result, RuntimeMessage> { + // SAFETY: all isolates stay on this executor; no scopes survive a turn. + unsafe { + self.isolate.enter(); + } + self.entered = true; + let result = match message { + HostMessage::Initialize { source } => self + .load(&source) + .map(|()| vec![RuntimeMessage::Initialized]), + HostMessage::Dispatch { + id, + export, + return_mode, + } => self.dispatch(id, &export, return_mode), + HostMessage::Close => Err(ErrorCode::InvalidProgram), + }; + let script_error = failures::take(&mut self.isolate); + // SAFETY: load/dispatch have released all scopes before yielding the executor. + unsafe { + self.isolate.exit(); + } + self.entered = false; + result.map_err(|code| RuntimeMessage::Error { code, script_error }) + } + + pub(super) fn load(&mut self, source: &str) -> Result<(), ErrorCode> { + if self.namespace.is_some() || source.len() > MAX_SOURCE { + return Err(ErrorCode::InvalidProgram); + } + v8::scope!(let scope, &mut self.isolate); + let context = v8::Local::new(scope, &self.context); + let scope = &mut v8::ContextScope::new(scope, context); + let tc = std::pin::pin!(v8::TryCatch::new(scope)); + let mut tc = tc.init(); + let result = (|| { + let source = v8::String::new(&tc, source).ok_or(ErrorCode::QuotaExceeded)?; + let name = v8::String::new(&tc, "plugin.mjs").ok_or(ErrorCode::QuotaExceeded)?; + let origin = v8::ScriptOrigin::new( + &tc, + name.into(), + /*resource_line_offset*/ 0, + /*resource_column_offset*/ 0, + /*resource_is_shared_cross_origin*/ false, + /*script_id*/ 0, + /*source_map_url*/ None, + /*resource_is_opaque*/ false, + /*is_wasm*/ false, + /*is_module*/ true, + /*host_defined_options*/ None, + ); + let mut source = v8::script_compiler::Source::new(source, Some(&origin)); + let module = v8::script_compiler::compile_module(&tc, &mut source) + .ok_or(ErrorCode::InvalidProgram)?; + // Reject imports before V8 invokes its platform-specific resolver callback. + if module.get_module_requests().length() != 0 { + let text = v8::String::new(&tc, "module imports are unavailable") + .ok_or(ErrorCode::QuotaExceeded)?; + let error = v8::Exception::type_error(&tc, text); + let message = v8::Exception::create_message(&tc, error); + failures::record(&mut tc, Some(message)); + return Err(ErrorCode::InvalidProgram); + } + module + .instantiate_module(&tc, |_, _, _, _| None) + .ok_or(ErrorCode::InvalidProgram)?; + let evaluation = module.evaluate(&tc).ok_or(ErrorCode::ExecutionFailed)?; + let evaluation = v8::Local::::try_from(evaluation) + .map_err(|_| ErrorCode::InvalidProgram)?; + failures::checkpoint(&mut tc)?; + if evaluation.state() != v8::PromiseState::Fulfilled { + return Err(ErrorCode::ExecutionFailed); + } + let namespace = v8::Local::::try_from(module.get_module_namespace()) + .map_err(|_| ErrorCode::InvalidProgram)?; + self.namespace = Some(v8::Global::new(&tc, namespace)); + Ok(()) + })(); + if result.is_err() { + let message = tc.message(); + failures::record(&mut tc, message); + } + result + } + + /// One JS turn; unresolved invocations yield back to the IPC dispatcher. + pub(super) fn dispatch( + &mut self, + id: u64, + export: &str, + return_mode: ReturnMode, + ) -> Result, ErrorCode> { + let namespace = self.namespace.as_ref().ok_or(ErrorCode::InvalidProgram)?; + v8::scope!(let scope, &mut self.isolate); + let context = v8::Local::new(scope, &self.context); + let scope = &mut v8::ContextScope::new(scope, context); + let tc = std::pin::pin!(v8::TryCatch::new(scope)); + let mut tc = tc.init(); + let result = (|| { + let mut replies = Vec::new(); + if export.is_empty() + || export.len() > MAX_METHOD + || self.pending_promises.contains_key(&id) + { + return Err(ErrorCode::QuotaExceeded); + } + let namespace = v8::Local::new(&tc, namespace); + let key = v8::String::new(&tc, export).ok_or(ErrorCode::QuotaExceeded)?; + let function = namespace + .get(&tc, key.into()) + .ok_or(ErrorCode::InvalidProgram)?; + let function = v8::Local::::try_from(function) + .map_err(|_| ErrorCode::InvalidProgram)?; + let receiver = v8::undefined(&tc).into(); + let result = function + .call(&tc, receiver, &[]) + .ok_or(ErrorCode::ExecutionFailed)?; + if return_mode == ReturnMode::MaybeAsync && result.is_promise() { + if self.pending_promises.len() >= MAX_PENDING_PROMISES { + return Err(ErrorCode::QuotaExceeded); + } + let promise = v8::Local::::try_from(result) + .map_err(|_| ErrorCode::InvalidResult)?; + self.pending_promises + .insert(id, v8::Global::new(&tc, promise)); + } else { + replies.push(RuntimeMessage::Reply { + id, + value: result_string(&mut tc, result)?, + }); + } + failures::checkpoint(&mut tc)?; + let completed: Vec<_> = self + .pending_promises + .extract_if(|_, promise| { + v8::Local::new(&tc, &*promise).state() != v8::PromiseState::Pending + }) + .collect(); + for (id, promise) in completed { + let promise = v8::Local::new(&tc, promise); + if promise.state() != v8::PromiseState::Fulfilled { + failures::record_rejection(&mut tc, promise); + return Err(ErrorCode::ExecutionFailed); + } + let value = promise.result(&tc); + replies.push(RuntimeMessage::Reply { + id, + value: result_string(&mut tc, value)?, + }); + } + if self.pending_promises.contains_key(&id) { + replies.push(RuntimeMessage::Pending { id }); + } + Ok(replies) + })(); + if result.is_err() { + let message = tc.message(); + failures::record(&mut tc, message); + } + result + } +} + +impl Drop for Engine { + fn drop(&mut self) { + // SAFETY: the executor drops only exited engines. OwnedIsolate's destructor exits it. + if !self.entered { + unsafe { + self.isolate.enter(); + } + } + } +} + +fn result_string( + scope: &mut v8::PinScope<'_, '_>, + result: v8::Local, +) -> Result { + let value = v8::Local::::try_from(result).map_err(|_| ErrorCode::InvalidResult)?; + if value.length() > MAX_RESULT || value.utf8_length(scope) > MAX_RESULT { + return Err(ErrorCode::QuotaExceeded); + } + Ok(value.to_rust_string_lossy(scope)) +} + +unsafe extern "C" fn reject_wasm( + _context: v8::Local, + _source: v8::Local, +) -> bool { + false +} diff --git a/codex-rs/code-mode-runtime/src/ui/failures.rs b/codex-rs/code-mode-runtime/src/ui/failures.rs new file mode 100644 index 0000000000..4a41f3631c --- /dev/null +++ b/codex-rs/code-mode-runtime/src/ui/failures.rs @@ -0,0 +1,99 @@ +//! Bounded per-isolate rejection tracking and script error capture. + +use codex_plugin_host::protocol::ErrorCode; +use codex_plugin_host::protocol::MAX_SCRIPT_ERROR_BYTES; +use codex_plugin_host::protocol::PluginScriptError; + +const MAX_REJECTIONS: usize = 32; + +#[derive(Default)] +struct Failures { + rejected: Vec>, + overflowed: bool, + script_error: Option, +} + +pub(super) fn install(isolate: &mut v8::OwnedIsolate) { + isolate.set_slot(Failures::default()); + isolate.set_promise_reject_callback(rejected); +} + +extern "C" fn rejected(message: v8::PromiseRejectMessage) { + // SAFETY: V8 supplies the callback's active isolate and scope. + v8::callback_scope!(unsafe scope, &message); + match message.get_event() { + v8::PromiseRejectEvent::PromiseRejectWithNoHandler => { + let promise = v8::Global::new(scope, message.get_promise()); + let Some(state) = scope.get_slot_mut::() else { + return; + }; + if state.rejected.len() == MAX_REJECTIONS { + state.overflowed = true; + } else { + state.rejected.push(promise); + } + } + v8::PromiseRejectEvent::PromiseHandlerAddedAfterReject => { + let promise = message.get_promise(); + if let Some(state) = scope.get_slot_mut::() { + state.rejected.retain(|candidate| candidate != &promise); + } + } + v8::PromiseRejectEvent::PromiseRejectAfterResolved + | v8::PromiseRejectEvent::PromiseResolveAfterResolved => {} + } +} + +pub(super) fn checkpoint(scope: &mut v8::PinScope<'_, '_>) -> Result<(), ErrorCode> { + scope.perform_microtask_checkpoint(); + if scope.is_execution_terminating() { + return Err(ErrorCode::ExecutionFailed); + } + let (rejected, overflowed) = { + let state = scope + .get_slot_mut::() + .ok_or(ErrorCode::ExecutionFailed)?; + ( + std::mem::take(&mut state.rejected), + std::mem::take(&mut state.overflowed), + ) + }; + if let Some(promise) = rejected.first() { + let promise = v8::Local::new(scope, promise); + record_rejection(scope, promise); + return Err(ErrorCode::ExecutionFailed); + } + if overflowed { + return Err(ErrorCode::QuotaExceeded); + } + Ok(()) +} + +pub(super) fn record_rejection(scope: &mut v8::PinScope<'_, '_>, promise: v8::Local) { + let reason = promise.result(scope); + let message = v8::Exception::create_message(scope, reason); + record(scope, Some(message)); +} + +pub(super) fn record(scope: &mut v8::PinScope<'_, '_>, message: Option>) { + let Some(message) = message else { return }; + let text = message.get(scope); + let full_len = text.utf8_length(scope); + let mut bytes = vec![0; full_len.min(MAX_SCRIPT_ERROR_BYTES)]; + let len = text.write_utf8_v2(scope, &mut bytes, v8::WriteFlags::kReplaceInvalidUtf8, None); + let script_error = PluginScriptError { + message: String::from_utf8_lossy(&bytes[..len]).into_owned(), + line: message.get_line_number(scope), + column: message.get_start_column().checked_add(1), + truncated: len < full_len, + }; + if let Some(state) = scope.get_slot_mut::() { + state.script_error = Some(script_error); + } +} + +pub(super) fn take(isolate: &mut v8::Isolate) -> Option { + isolate + .get_slot_mut::() + .and_then(|state| state.script_error.take()) +} diff --git a/codex-rs/code-mode-runtime/src/ui/mod.rs b/codex-rs/code-mode-runtime/src/ui/mod.rs new file mode 100644 index 0000000000..603e2cb287 --- /dev/null +++ b/codex-rs/code-mode-runtime/src/ui/mod.rs @@ -0,0 +1,23 @@ +//! Lightweight UI isolates share one bounded executor in the plugin process. +//! No ordinary CodeMode tools or globals. +//! Script failures discard their isolate without affecting its peers. + +mod engine; +mod failures; +mod transport; +mod watchdog; +mod workers; + +use std::io; + +pub fn run() -> io::Result<()> { + crate::initialize_v8(crate::V8JitMode::Disabled).map_err(io::Error::other)?; + unsafe extern "C" { + fn v8__V8__IsSandboxEnabled() -> bool; + } + // SAFETY: linked V8 exposes this query and has already been initialized. + if !unsafe { v8__V8__IsSandboxEnabled() } { + return Err(io::Error::other("sandbox-enabled V8 required")); + } + workers::run() +} diff --git a/codex-rs/code-mode-runtime/src/ui/transport.rs b/codex-rs/code-mode-runtime/src/ui/transport.rs new file mode 100644 index 0000000000..f4893de704 --- /dev/null +++ b/codex-rs/code-mode-runtime/src/ui/transport.rs @@ -0,0 +1,133 @@ +//! Framed I/O and per-instance queue admission for the plugin child. + +use codex_plugin_host::framing; +use codex_plugin_host::protocol::MAX_IN_FLIGHT_TURNS; +use codex_plugin_host::protocol::Packet; +use codex_plugin_host::protocol::RuntimeMessage; +use futures::io::AllowStdIo; +use std::io; +use std::sync::Arc; +use std::sync::OnceLock; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::sync::mpsc; +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::Semaphore; +use tokio_util::codec::FramedRead; +use tokio_util::codec::FramedWrite; +use tokio_util::codec::LengthDelimitedCodec; +use tokio_util::compat::Compat; +use tokio_util::compat::FuturesAsyncReadCompatExt; +use tokio_util::compat::FuturesAsyncWriteCompatExt; + +// A reply may reach the parent just before its writer releases the output slot. +pub(super) const OUTPUT_SLOTS_PER_INSTANCE: usize = MAX_IN_FLIGHT_TURNS + 1; + +/// Only poll on a dedicated blocking thread: AllowStdIo blocks inside poll_read. +pub(super) fn blocking_reader( + input: R, +) -> FramedRead>, LengthDelimitedCodec> { + FramedRead::new(AllowStdIo::new(input).compat(), framing::codec()) +} + +/// Only poll on a dedicated blocking thread: AllowStdIo blocks inside poll_write. +pub(super) fn blocking_writer( + output: W, +) -> FramedWrite>, LengthDelimitedCodec> { + FramedWrite::new(AllowStdIo::new(output).compat_write(), framing::codec()) +} + +pub(super) struct Instance { + pub(super) id: u64, + pub(super) retired: Arc, + pub(super) close_output: OnceLock<(mpsc::SyncSender, QueuedOutput)>, + pub(super) closed: AtomicBool, + pub(super) input_slots: Arc, + output_slots: Arc, +} + +impl Instance { + pub(super) fn new(id: u64) -> Self { + Self { + id, + retired: Arc::new(AtomicBool::new(false)), + close_output: OnceLock::new(), + closed: AtomicBool::new(false), + input_slots: Arc::new(Semaphore::new(MAX_IN_FLIGHT_TURNS)), + output_slots: Arc::new(Semaphore::new(OUTPUT_SLOTS_PER_INSTANCE)), + } + } +} + +impl Drop for Instance { + fn drop(&mut self) { + // The engine, queued turns and earlier output have all released this instance. + if let Some((outgoing_tx, output)) = self.close_output.take() + && let Err(error) = outgoing_tx.try_send(output) + { + super::workers::fatal(io::Error::other(error)); + } + } +} + +pub(super) enum OutputReservation { + Turn { + _instance: Arc, + _permit: OwnedSemaphorePermit, + }, + Close { + retired: Arc, + }, + Rejected { + _permit: OwnedSemaphorePermit, + }, +} + +pub(super) struct QueuedOutput { + pub(super) packets: Vec>, + pub(super) _reservation: OutputReservation, +} + +pub(super) struct Transport { + pub(super) instance: Arc, + pub(super) outgoing_tx: mpsc::SyncSender, +} + +impl Transport { + pub(super) fn ensure_open(&self) -> io::Result<()> { + if self.instance.closed.load(Ordering::Acquire) { + return Err(io::Error::other("plugin instance closed")); + } + Ok(()) + } + + pub(super) fn write(&self, messages: Vec) -> io::Result<()> { + self.ensure_open()?; + if messages.is_empty() { + return Ok(()); + } + let permit = Arc::clone(&self.instance.output_slots) + .try_acquire_owned() + .map_err(|_| io::Error::other("plugin instance output full"))?; + let packets = messages + .into_iter() + .map(|message| Packet { + instance: self.instance.id, + message, + }) + .collect(); + self.outgoing_tx + .try_send(QueuedOutput { + packets, + _reservation: OutputReservation::Turn { + _instance: Arc::clone(&self.instance), + _permit: permit, + }, + }) + .map_err(io::Error::other) + } +} + +#[cfg(test)] +#[path = "transport_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-runtime/src/ui/transport_tests.rs b/codex-rs/code-mode-runtime/src/ui/transport_tests.rs new file mode 100644 index 0000000000..c2e9f19319 --- /dev/null +++ b/codex-rs/code-mode-runtime/src/ui/transport_tests.rs @@ -0,0 +1,71 @@ +#![cfg(test)] +//! Native sends fail immediately on saturation or cancellation. + +use super::*; +use futures::SinkExt; +use futures::StreamExt; +use futures::executor::block_on; +use pretty_assertions::assert_eq; + +#[test] +fn a_full_instance_output_does_not_block_peers_or_the_executor() { + let (outgoing_tx, outgoing_rx) = mpsc::sync_channel(OUTPUT_SLOTS_PER_INSTANCE + 1); + let first = Transport { + instance: Arc::new(Instance::new(/*id*/ 1)), + outgoing_tx: outgoing_tx.clone(), + }; + let peer = Transport { + instance: Arc::new(Instance::new(/*id*/ 2)), + outgoing_tx, + }; + for _ in 0..OUTPUT_SLOTS_PER_INSTANCE { + first.write(vec![RuntimeMessage::Initialized]).unwrap(); + } + assert!(first.write(vec![RuntimeMessage::Initialized]).is_err()); + peer.write(vec![RuntimeMessage::Initialized]).unwrap(); + assert_eq!(outgoing_rx.recv().unwrap().packets[0].instance, 1); + first.write(vec![RuntimeMessage::Initialized]).unwrap(); + first.instance.closed.store(true, Ordering::Release); + assert!(first.write(vec![RuntimeMessage::Initialized]).is_err()); +} + +#[test] +fn rejects_length_before_reading_payload() { + let bytes = u32::MAX.to_le_bytes(); + assert_eq!( + block_on(blocking_reader(bytes.as_slice()).next()) + .unwrap() + .unwrap_err() + .kind(), + io::ErrorKind::InvalidData + ); +} + +#[test] +fn rejects_truncated_payload() { + assert_eq!( + block_on(blocking_reader([1, 0, 0, 0].as_slice()).next()) + .unwrap_or_else(|| Err(io::ErrorKind::UnexpectedEof.into())) + .unwrap_err() + .kind(), + io::ErrorKind::UnexpectedEof + ); +} + +#[test] +fn blocking_transport_preserves_wire_format_and_buffered_frames() -> io::Result<()> { + let payloads = [b"abc".as_slice(), b"", b"hi\0"]; + let mut bytes = Vec::new(); + { + let mut output = blocking_writer(&mut bytes); + for payload in payloads { + block_on(output.send(payload.to_vec().into()))?; + } + } + assert_eq!(bytes, b"\x03\0\0\0abc\0\0\0\0\x03\0\0\0hi\0"); + let mut input = blocking_reader(bytes.as_slice()); + for payload in payloads { + assert_eq!(block_on(input.next()).unwrap()?.as_ref(), payload); + } + Ok(()) +} diff --git a/codex-rs/code-mode-runtime/src/ui/watchdog.rs b/codex-rs/code-mode-runtime/src/ui/watchdog.rs new file mode 100644 index 0000000000..30e19ae438 --- /dev/null +++ b/codex-rs/code-mode-runtime/src/ui/watchdog.rs @@ -0,0 +1,143 @@ +//! One observer interrupts the active JS turn; disarming and interruption share a lock. + +use std::io; +use std::sync::Arc; +use std::sync::Condvar; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::thread::JoinHandle; +use std::time::Duration; +use std::time::Instant; + +struct Active { + instance: u64, + handle: v8::IsolateHandle, + deadline: Instant, + terminated: bool, +} + +#[derive(Default)] +struct Execution { + active: Option, + stopped: bool, +} + +type Shared = Arc<(Mutex, Condvar)>; + +#[derive(Clone)] +pub(super) struct Observer(Shared); + +pub(super) struct IsolateExecutionWatchdog { + pub(super) observer: Observer, + thread: Option>, +} + +impl IsolateExecutionWatchdog { + pub(super) fn new() -> io::Result { + let shared = Arc::new((Mutex::new(Execution::default()), Condvar::new())); + let observer = Observer(Arc::clone(&shared)); + let thread = std::thread::Builder::new() + .name("ui-plugin-watchdog".into()) + .spawn(move || { + let (lock, changed) = &*shared; + let mut state = lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + while !state.stopped { + if let Some(active) = &mut state.active + && !active.terminated + { + let remaining = active.deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + active.handle.terminate_execution(); + active.terminated = true; + } else { + (state, _) = changed + .wait_timeout(state, remaining) + .unwrap_or_else(std::sync::PoisonError::into_inner); + continue; + } + } + state = changed + .wait(state) + .unwrap_or_else(std::sync::PoisonError::into_inner); + } + })?; + Ok(Self { + observer, + thread: Some(thread), + }) + } + + pub(super) fn arm(&self, instance: u64, handle: v8::IsolateHandle) { + let (lock, changed) = &*self.observer.0; + let mut state = lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert!(state.active.is_none()); + state.active = Some(Active { + instance, + handle, + deadline: Instant::now() + Duration::from_millis(100), + terminated: false, + }); + changed.notify_one(); + } + + pub(super) fn disarm(&self) -> io::Result<()> { + self.disarm_with(|lock| { + lock.lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + }) + } + + fn disarm_with<'a>( + &'a self, + acquire: impl FnOnce(&'a Mutex) -> MutexGuard<'a, Execution>, + ) -> io::Result<()> { + let (lock, changed) = &*self.observer.0; + let active = acquire(lock).active.take(); + changed.notify_one(); + if active.is_some_and(|active| active.terminated) { + Err(io::Error::new( + io::ErrorKind::TimedOut, + "plugin execution deadline exceeded", + )) + } else { + Ok(()) + } + } +} + +#[cfg(test)] +#[path = "watchdog_tests.rs"] +mod tests; + +impl Observer { + pub(super) fn cancel(&self, instance: u64) { + let (lock, changed) = &*self.0; + let mut state = lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(active) = &mut state.active + && active.instance == instance + { + active.handle.terminate_execution(); + active.terminated = true; + } + changed.notify_one(); + } +} + +impl Drop for IsolateExecutionWatchdog { + fn drop(&mut self) { + let (lock, changed) = &*self.observer.0; + lock.lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .stopped = true; + changed.notify_one(); + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } +} diff --git a/codex-rs/code-mode-runtime/src/ui/watchdog_tests.rs b/codex-rs/code-mode-runtime/src/ui/watchdog_tests.rs new file mode 100644 index 0000000000..64da2a34be --- /dev/null +++ b/codex-rs/code-mode-runtime/src/ui/watchdog_tests.rs @@ -0,0 +1,38 @@ +#![cfg(test)] +use super::*; +use pretty_assertions::assert_eq; +use std::sync::mpsc; + +#[test] +fn disarm_waits_for_an_interrupt_in_progress() { + crate::v8_init::ensure_v8_initialized().unwrap(); + let isolate = v8::Isolate::new(Default::default()); + let watchdog = IsolateExecutionWatchdog { + observer: Observer(Arc::new((Mutex::new(Execution::default()), Condvar::new()))), + thread: None, + }; + watchdog.arm(/*instance*/ 1, isolate.thread_safe_handle()); + let (lock, _) = &*watchdog.observer.0; + let mut interrupt = lock.lock().unwrap(); + std::thread::scope(|scope| { + let (attempt_tx, attempt_rx) = mpsc::sync_channel(0); + let watchdog = &watchdog; + let disarming = scope.spawn(move || { + watchdog.disarm_with(|lock| { + assert!(matches!( + lock.try_lock(), + Err(std::sync::TryLockError::WouldBlock) + )); + attempt_tx.send(()).unwrap(); + lock.lock().unwrap() + }) + }); + attempt_rx.recv().unwrap(); + interrupt.active.as_mut().unwrap().terminated = true; + drop(interrupt); + assert_eq!( + disarming.join().unwrap().unwrap_err().kind(), + io::ErrorKind::TimedOut + ); + }); +} diff --git a/codex-rs/code-mode-runtime/src/ui/workers.rs b/codex-rs/code-mode-runtime/src/ui/workers.rs new file mode 100644 index 0000000000..d2636e0cb7 --- /dev/null +++ b/codex-rs/code-mode-runtime/src/ui/workers.rs @@ -0,0 +1,286 @@ +//! Dedicated stdin, stdout, JS execution and watchdog threads. Queue admission is +//! per instance, and cancellation bypasses queued JS turns; V8 manages its own workers. + +use std::collections::HashMap; +use std::io; +use std::sync::Arc; +use std::sync::Weak; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::sync::mpsc; +use std::thread; + +use super::engine::Engine; +use super::transport::Instance; +use super::transport::OUTPUT_SLOTS_PER_INSTANCE; +use super::transport::OutputReservation; +use super::transport::QueuedOutput; +use super::transport::Transport; +use super::transport::blocking_reader; +use super::transport::blocking_writer; +use super::watchdog::IsolateExecutionWatchdog; +use codex_plugin_host::MAX_INSTANCES; +use codex_plugin_host::protocol::ErrorCode; +use codex_plugin_host::protocol::HostMessage; +use codex_plugin_host::protocol::MAX_IN_FLIGHT_TURNS; +use codex_plugin_host::protocol::Packet; +use codex_plugin_host::protocol::RuntimeMessage; +use futures::SinkExt; +use futures::StreamExt; +use futures::executor::block_on; +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::Semaphore; + +// Every instance reserves space for its turns and one close; output also reserves +// one slot for a rejected open. +const COMMAND_QUEUE_CAPACITY: usize = MAX_INSTANCES * (MAX_IN_FLIGHT_TURNS + 1); +const OUTPUT_QUEUE_CAPACITY: usize = MAX_INSTANCES * (OUTPUT_SLOTS_PER_INSTANCE + 1) + 1; + +enum Command { + Turn { + packet: Packet, + instance: Arc, + _permit: OwnedSemaphorePermit, + }, + Close(Arc), +} + +pub(super) fn run() -> io::Result<()> { + let (outgoing_tx, outgoing_rx) = mpsc::sync_channel::(OUTPUT_QUEUE_CAPACITY); + thread::Builder::new() + .name("ui-plugin-output".into()) + .spawn(move || { + let mut writer = blocking_writer(io::stdout().lock()); + if let Err(error) = block_on(writer.send(codex_plugin_host::READY.to_vec().into())) { + fatal(error); + } + for output in outgoing_rx { + // The parent can send a replacement as soon as it reads Closed. + if let OutputReservation::Close { retired } = &output._reservation { + retired.store(true, Ordering::Release); + } + for packet in &output.packets { + let bytes = + serde_json::to_vec(packet).unwrap_or_else(|error| fatal(error.into())); + if let Err(error) = block_on(writer.send(bytes.into())) { + fatal(error); + } + } + } + })?; + let watchdog = IsolateExecutionWatchdog::new()?; + let observer = watchdog.observer.clone(); + let (command_tx, command_rx) = mpsc::sync_channel(COMMAND_QUEUE_CAPACITY); + let executor_outgoing_tx = outgoing_tx.clone(); + thread::Builder::new() + .name("ui-plugin-executor".into()) + .spawn(move || { + if let Err(error) = execute(command_rx, executor_outgoing_tx, watchdog) { + fatal(error); + } + })?; + let mut live = HashMap::, Weak)>::new(); + let rejected_slot = Arc::new(Semaphore::new(1)); + let mut last_id = None; + let mut input = blocking_reader(io::stdin().lock()); + loop { + let packet: Packet = + serde_json::from_slice(&block_on(input.next()).ok_or(io::ErrorKind::UnexpectedEof)??)?; + // An instance keeps its registration until its engine, queued turns and output drain. + live.retain(|_, (retired, _)| !retired.load(Ordering::Acquire)); + let instance = match &packet.message { + HostMessage::Initialize { .. } => { + if last_id.is_some_and(|last| packet.instance <= last) { + return Err(io::Error::other("invalid plugin instance ID")); + } + last_id = Some(packet.instance); + if live.len() >= MAX_INSTANCES { + // Backpressure excess opens instead of accumulating unbounded close frames. + let permit = block_on(Arc::clone(&rejected_slot).acquire_owned()) + .map_err(io::Error::other)?; + outgoing_tx + .try_send(QueuedOutput { + packets: vec![close_packet(packet.instance)], + _reservation: OutputReservation::Rejected { _permit: permit }, + }) + .map_err(io::Error::other)?; + continue; + } + let instance = Arc::new(Instance::new(packet.instance)); + live.insert( + instance.id, + (Arc::clone(&instance.retired), Arc::downgrade(&instance)), + ); + instance + } + HostMessage::Dispatch { .. } | HostMessage::Close => { + let Some(instance) = live + .get(&packet.instance) + .and_then(|(_, instance)| instance.upgrade()) + else { + continue; + }; + instance + } + }; + admit_turn(packet, instance, &observer, &command_tx)?; + } +} + +fn admit_turn( + packet: Packet, + instance: Arc, + observer: &super::watchdog::Observer, + command_tx: &mpsc::SyncSender, +) -> io::Result<()> { + if instance.closed.load(Ordering::Acquire) { + return Ok(()); + } + let permit = if matches!(packet.message, HostMessage::Close) { + None + } else { + Arc::clone(&instance.input_slots).try_acquire_owned().ok() + }; + let command = if let Some(permit) = permit { + Command::Turn { + packet, + instance, + _permit: permit, + } + } else { + if instance.closed.swap(true, Ordering::AcqRel) { + return Ok(()); + } + observer.cancel(instance.id); + Command::Close(instance) + }; + command_tx.try_send(command).map_err(io::Error::other) +} + +fn execute( + command_rx: mpsc::Receiver, + outgoing_tx: mpsc::SyncSender, + watchdog: IsolateExecutionWatchdog, +) -> io::Result<()> { + let mut instances = HashMap::::new(); + for command in command_rx { + let (packet, instance, _permit) = match command { + Command::Close(instance) => { + if instances.remove(&instance.id).is_some() { + send_close(&outgoing_tx, instance, /*error*/ None)?; + } + continue; + } + Command::Turn { + packet, + instance, + _permit, + } => (packet, instance, _permit), + }; + if instance.closed.load(Ordering::Acquire) { + if matches!(packet.message, HostMessage::Initialize { .. }) + || instances.remove(&instance.id).is_some() + { + send_close(&outgoing_tx, instance, /*error*/ None)?; + } + continue; + } + if matches!(packet.message, HostMessage::Initialize { .. }) { + let transport = Transport { + instance: Arc::clone(&instance), + outgoing_tx: outgoing_tx.clone(), + }; + match Engine::new(transport) { + Ok(engine) => { + instances.insert(instance.id, engine); + } + Err(error) => { + eprintln!("plugin isolate initialization failed: {error}"); + instance.closed.store(true, Ordering::Release); + send_close(&outgoing_tx, instance, /*error*/ None)?; + continue; + } + } + } + let Some(engine) = instances.get_mut(&instance.id) else { + let was_closed = instance.closed.swap(true, Ordering::AcqRel); + let error = (!was_closed).then_some(RuntimeMessage::Error { + code: ErrorCode::ExecutionFailed, + script_error: None, + }); + send_close(&outgoing_tx, instance, error)?; + continue; + }; + watchdog.arm(instance.id, engine.handle()); + let result = engine + .transport + .ensure_open() + .map(|()| engine.turn(packet.message)); + drop(_permit); + let error = match (watchdog.disarm(), result) { + (Ok(()), Ok(Ok(messages))) => { + if engine.transport.write(messages).is_ok() { + continue; + } + None + } + (Ok(()), Ok(Err(error))) => Some(error), + (Ok(()), Err(_)) => None, + (Err(_), _) => Some(RuntimeMessage::Error { + code: ErrorCode::ExecutionTimedOut, + script_error: None, + }), + }; + let error = if instance.closed.swap(true, Ordering::AcqRel) { + None + } else { + error + }; + instances.remove(&instance.id); + send_close(&outgoing_tx, instance, error)?; + } + Err(io::Error::other("plugin input stopped")) +} + +fn close_packet(instance: u64) -> Packet { + Packet { + instance, + message: RuntimeMessage::Closed, + } +} + +fn send_close( + outgoing_tx: &mpsc::SyncSender, + instance: Arc, + error: Option, +) -> io::Result<()> { + let mut packets = Vec::new(); + if let Some(error) = error { + packets.push(Packet { + instance: instance.id, + message: error, + }); + } + packets.push(close_packet(instance.id)); + instance + .close_output + .set(( + outgoing_tx.clone(), + QueuedOutput { + packets, + _reservation: OutputReservation::Close { + retired: Arc::clone(&instance.retired), + }, + }, + )) + .map_err(|_| io::Error::other("plugin close already scheduled")) +} + +pub(super) fn fatal(error: io::Error) -> ! { + eprintln!("plugin helper failed: {error}"); + std::process::exit(1); +} + +#[cfg(test)] +#[path = "workers_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-runtime/src/ui/workers_tests.rs b/codex-rs/code-mode-runtime/src/ui/workers_tests.rs new file mode 100644 index 0000000000..5c66dd1990 --- /dev/null +++ b/codex-rs/code-mode-runtime/src/ui/workers_tests.rs @@ -0,0 +1,126 @@ +#![cfg(test)] +//! Exercise executor invariants that the stdin reader normally enforces. + +use super::*; +use codex_plugin_host::protocol::ReturnMode; +use pretty_assertions::assert_eq; +use serde_json::json; + +#[test] +fn missing_engine_fails_only_its_instance() { + crate::v8_init::ensure_v8_initialized().unwrap(); + let peer = Arc::new(Instance::new(/*id*/ 1)); + let missing = Arc::new(Instance::new(/*id*/ 2)); + let dispatch = || HostMessage::Dispatch { + id: 0, + export: "default".into(), + return_mode: ReturnMode::Sync, + }; + let (command_tx, command_rx) = mpsc::channel(); + for (instance, message) in [ + ( + &peer, + HostMessage::Initialize { + source: "export default () => 'alive'".into(), + }, + ), + (&missing, dispatch()), + (&peer, dispatch()), + ] { + command_tx + .send(Command::Turn { + packet: Packet { + instance: instance.id, + message, + }, + instance: Arc::clone(instance), + _permit: Arc::clone(&instance.input_slots) + .try_acquire_owned() + .unwrap(), + }) + .unwrap(); + } + drop(command_tx); + let (outgoing_tx, outgoing_rx) = mpsc::sync_channel(OUTPUT_QUEUE_CAPACITY); + let error = execute( + command_rx, + outgoing_tx, + IsolateExecutionWatchdog::new().unwrap(), + ) + .unwrap_err(); + assert_eq!(error.to_string(), "plugin input stopped"); + assert!(missing.closed.load(Ordering::Acquire)); + drop(missing); + let packets: Vec<_> = outgoing_rx + .into_iter() + .flat_map(|output| output.packets) + .collect(); + assert_eq!( + serde_json::to_value(packets).unwrap(), + json!([ + {"instance":1,"message":{"type":"initialized"}}, + {"instance":1,"message":{"type":"reply","id":0,"value":"alive"}}, + {"instance":2,"message":{"type":"error","code":"execution_failed","scriptError":null}}, + {"instance":2,"message":{"type":"closed"}} + ]) + ); +} + +#[test] +fn close_waits_for_the_last_queued_turn() { + let (outgoing_tx, outgoing_rx) = mpsc::sync_channel(OUTPUT_QUEUE_CAPACITY); + let instance = Arc::new(Instance::new(/*id*/ 1)); + let queued_turn = Arc::clone(&instance); + send_close(&outgoing_tx, instance, /*error*/ None).unwrap(); + assert!(matches!( + outgoing_rx.try_recv(), + Err(mpsc::TryRecvError::Empty) + )); + drop(queued_turn); + let output = outgoing_rx.try_recv().unwrap(); + assert_eq!( + serde_json::to_value(output.packets).unwrap(), + json!([{"instance":1,"message":{"type":"closed"}}]) + ); +} + +#[test] +fn input_overload_closes_only_the_busy_instance() { + let busy = Arc::new(Instance::new(/*id*/ 1)); + let peer = Arc::new(Instance::new(/*id*/ 2)); + let watchdog = IsolateExecutionWatchdog::new().unwrap(); + let (command_tx, command_rx) = mpsc::sync_channel(COMMAND_QUEUE_CAPACITY); + let dispatch = |instance: &Arc, id| { + let packet = Packet { + instance: instance.id, + message: HostMessage::Dispatch { + id, + export: "run".into(), + return_mode: ReturnMode::Sync, + }, + }; + admit_turn( + packet, + Arc::clone(instance), + &watchdog.observer, + &command_tx, + ) + .unwrap(); + }; + // Leave every turn queued so its permit stays in flight. + for id in 0..=MAX_IN_FLIGHT_TURNS as u64 { + dispatch(&busy, id); + } + dispatch(&peer, 0); + for _ in 0..MAX_IN_FLIGHT_TURNS { + assert!( + matches!(command_rx.try_recv(), Ok(Command::Turn { instance, .. }) if instance.id == busy.id) + ); + } + assert!( + matches!(command_rx.try_recv(), Ok(Command::Close(instance)) if instance.id == busy.id) + ); + assert!( + matches!(command_rx.try_recv(), Ok(Command::Turn { instance, .. }) if instance.id == peer.id) + ); +} diff --git a/codex-rs/plugin-host/Cargo.toml b/codex-rs/plugin-host/Cargo.toml index 31e56dffb2..98e88b4262 100644 --- a/codex-rs/plugin-host/Cargo.toml +++ b/codex-rs/plugin-host/Cargo.toml @@ -29,6 +29,7 @@ tokio-util = { workspace = true, features = ["codec", "compat", "rt"] } codex-utils-process = { workspace = true } codex-utils-pty = { workspace = true } libc = { workspace = true } +serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } tokio = { workspace = true, features = ["io-util", "macros", "process", "rt", "sync", "time"] } diff --git a/codex-rs/plugin-host/src/framing.rs b/codex-rs/plugin-host/src/framing.rs index 24e2e87e93..a90257d0cf 100644 --- a/codex-rs/plugin-host/src/framing.rs +++ b/codex-rs/plugin-host/src/framing.rs @@ -2,7 +2,7 @@ use tokio_util::codec::LengthDelimitedCodec; -pub const MAX_FRAME: usize = 256 * 1024; +pub const MAX_FRAME: usize = 1024 * 1024; pub fn codec() -> LengthDelimitedCodec { LengthDelimitedCodec::builder() diff --git a/codex-rs/plugin-host/src/lib.rs b/codex-rs/plugin-host/src/lib.rs index 677f24d03b..787adf7ccf 100644 --- a/codex-rs/plugin-host/src/lib.rs +++ b/codex-rs/plugin-host/src/lib.rs @@ -2,10 +2,13 @@ pub mod framing; mod process; +pub mod protocol; pub use process::PluginExit; pub use process::PluginProcess; pub use process::PluginProcessError; pub const HELPER_ARGUMENT: &str = "--codex-ui-plugin-process"; -pub const READY: &[u8] = b"codex-ui-plugin-process/2"; +pub const READY: &[u8] = b"codex-ui-plugin-process/3"; + +pub const MAX_INSTANCES: usize = 32; diff --git a/codex-rs/plugin-host/src/protocol.rs b/codex-rs/plugin-host/src/protocol.rs new file mode 100644 index 0000000000..3ce0970a8f --- /dev/null +++ b/codex-rs/plugin-host/src/protocol.rs @@ -0,0 +1,88 @@ +//! Narrow experimental ABI. The child supplies correlation IDs, never authority or thread IDs. + +use serde::Deserialize; +use serde::Serialize; + +// Leave room for worst-case JSON escaping and the packet envelope within MAX_FRAME. +pub const MAX_SOURCE: usize = 128 * 1024; +pub const MAX_RESULT: usize = 64 * 1024; +pub const MAX_METHOD: usize = 1024; +pub const MAX_PENDING_PROMISES: usize = 1024; +pub const MAX_IN_FLIGHT_TURNS: usize = 8; +pub const MAX_SCRIPT_ERROR_BYTES: usize = 16 * 1024; +const _: () = assert!(MAX_SOURCE * 6 + 1024 <= crate::framing::MAX_FRAME); + +#[derive(Debug, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Packet { + pub instance: u64, + pub message: T, +} + +#[derive(Debug, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "camelCase", deny_unknown_fields)] +pub enum HostMessage { + Initialize { + source: String, + }, + Dispatch { + id: u64, + export: String, + #[serde(rename = "returnMode")] + return_mode: ReturnMode, + }, + Close, +} + +#[derive(Debug, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "camelCase", deny_unknown_fields)] +pub enum RuntimeMessage { + Initialized, + Pending { + id: u64, + }, + Reply { + id: u64, + value: String, + }, + Error { + code: ErrorCode, + #[serde(rename = "scriptError")] + script_error: Option, + }, + Closed, +} + +/// Original script error text and source location. Callers must sanitize text for display. +#[derive(Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PluginScriptError { + pub message: String, + pub line: Option, + pub column: Option, + pub truncated: bool, +} + +impl std::fmt::Debug for PluginScriptError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("PluginScriptError([redacted])") + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum ReturnMode { + Sync, + MaybeAsync, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ErrorCode { + InvalidProgram, + InvalidResult, + ExecutionFailed, + ExecutionTimedOut, + QuotaExceeded, + InstanceClosed, +}