feat(code-mode): allow disabling V8 JIT (#31303)

We want the option to be able to run code-mode in jitless mode.

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

@@ -3,6 +3,7 @@ mod remote_session;
 mod runtime;
 mod service;
 mod session_runtime;
+mod v8_init;
 
 pub(crate) type TaskFailureHandler = std::sync::Arc<dyn Fn(String) + Send + Sync>;
 
@@ -12,3 +13,5 @@ pub use remote_session::ProcessOwnedCodeModeSessionProvider;
 pub use service::InProcessCodeModeSession;
 pub use service::InProcessCodeModeSessionProvider;
 pub use service::NoopCodeModeSessionDelegate;
+pub use v8_init::V8JitMode;
+pub use v8_init::initialize_v8;

codex-rs/code-mode/src/runtime/mod.rs

@@ -7,7 +7,6 @@ mod value;
 use std::collections::HashMap;
 use std::panic::AssertUnwindSafe;
 use std::panic::catch_unwind;
-use std::sync::OnceLock;
 use std::sync::mpsc as std_mpsc;
 use std::thread;
 
@@ -21,6 +20,7 @@ use serde_json::Value as JsonValue;
 use tokio::sync::mpsc;
 
 use crate::TaskFailureHandler;
+use crate::v8_init::ensure_v8_initialized;
 
 const EXIT_SENTINEL: &str = "__codex_code_mode_exit__";
 
@@ -84,7 +84,7 @@ pub(crate) fn spawn_runtime(
     ),
     String,
 > {
-    initialize_v8()?;
+    ensure_v8_initialized()?;
 
     let (command_tx, command_rx) = std_mpsc::channel();
     let (control_tx, control_rx) = std_mpsc::channel();
@@ -165,22 +165,6 @@ pub(super) enum CompletionState {
     },
 }
 
-fn initialize_v8() -> Result<(), String> {
-    static PLATFORM: OnceLock<Result<v8::SharedRef<v8::Platform>, String>> = OnceLock::new();
-
-    match PLATFORM.get_or_init(|| {
-        v8::icu::set_common_data_77(deno_core_icudata::ICU_DATA)
-            .map_err(|error_code| format!("failed to initialize ICU data: {error_code}"))?;
-        let platform = v8::new_default_platform(0, false).make_shared();
-        v8::V8::initialize_platform(platform.clone());
-        v8::V8::initialize();
-        Ok(platform)
-    }) {
-        Ok(_) => Ok(()),
-        Err(error_text) => Err(error_text.clone()),
-    }
-}
-
 fn run_runtime(
     config: RuntimeConfig,
     event_tx: mpsc::UnboundedSender<RuntimeEvent>,

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

@@ -0,0 +1,65 @@
+use std::sync::OnceLock;
+
+/// Controls whether V8 may generate executable code at runtime.
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
+pub enum V8JitMode {
+    #[default]
+    Enabled,
+    Disabled,
+}
+
+struct V8Initialization {
+    _platform: v8::SharedRef<v8::Platform>,
+    jit_mode: V8JitMode,
+}
+
+static V8_INITIALIZATION: OnceLock<Result<V8Initialization, String>> = OnceLock::new();
+
+/// Initializes the process-wide V8 platform with the requested JIT mode.
+///
+/// Call this before executing any code-mode cells when JIT must be disabled.
+/// V8 cannot change JIT mode after initialization, so a later call requesting
+/// a different mode returns an error. Code mode initializes V8 with JIT enabled
+/// by default when this function has not been called explicitly.
+pub fn initialize_v8(jit_mode: V8JitMode) -> Result<(), String> {
+    match V8_INITIALIZATION.get_or_init(|| initialize_v8_with_mode(jit_mode)) {
+        Ok(initialization) if initialization.jit_mode == jit_mode => Ok(()),
+        Ok(initialization) => Err(format!(
+            "V8 was already initialized with JIT {}",
+            initialization.jit_mode.description()
+        )),
+        Err(error_text) => Err(error_text.clone()),
+    }
+}
+
+pub(crate) fn ensure_v8_initialized() -> Result<(), String> {
+    match V8_INITIALIZATION.get_or_init(|| initialize_v8_with_mode(V8JitMode::Enabled)) {
+        Ok(_) => Ok(()),
+        Err(error_text) => Err(error_text.clone()),
+    }
+}
+
+fn initialize_v8_with_mode(jit_mode: V8JitMode) -> Result<V8Initialization, String> {
+    v8::icu::set_common_data_77(deno_core_icudata::ICU_DATA)
+        .map_err(|error_code| format!("failed to initialize ICU data: {error_code}"))?;
+    match jit_mode {
+        V8JitMode::Enabled => {}
+        V8JitMode::Disabled => v8::V8::set_flags_from_string("--jitless"),
+    }
+    let platform = v8::new_default_platform(0, false).make_shared();
+    v8::V8::initialize_platform(platform.clone());
+    v8::V8::initialize();
+    Ok(V8Initialization {
+        _platform: platform,
+        jit_mode,
+    })
+}
+
+impl V8JitMode {
+    fn description(self) -> &'static str {
+        match self {
+            Self::Enabled => "enabled",
+            Self::Disabled => "disabled",
+        }
+    }
+}

codex-rs/code-mode/tests/jit.rs

@@ -0,0 +1,41 @@
+use codex_code_mode::ExecuteRequest;
+use codex_code_mode::InProcessCodeModeSession;
+use codex_code_mode::RuntimeResponse;
+use codex_code_mode::V8JitMode;
+use codex_code_mode::initialize_v8;
+use pretty_assertions::assert_eq;
+
+#[tokio::test]
+async fn code_mode_runs_with_jit_disabled() {
+    initialize_v8(V8JitMode::Disabled).expect("initialize V8 without JIT");
+
+    let service = InProcessCodeModeSession::new();
+    let started = service
+        .execute(ExecuteRequest {
+            tool_call_id: "call_1".to_string(),
+            enabled_tools: Vec::new(),
+            source: "21 * 2;".to_string(),
+            yield_time_ms: None,
+            max_output_tokens: None,
+        })
+        .await
+        .expect("start code-mode cell");
+    let cell_id = started.cell_id.clone();
+    let response = started
+        .initial_response()
+        .await
+        .expect("execute code-mode cell");
+
+    assert_eq!(
+        response,
+        RuntimeResponse::Result {
+            cell_id,
+            content_items: Vec::new(),
+            error_text: None,
+        }
+    );
+    assert_eq!(
+        initialize_v8(V8JitMode::Enabled),
+        Err("V8 was already initialized with JIT disabled".to_string())
+    );
+}