Prewarm code-mode host connections at session startup (#40678)

## What changed

- Add the opt-in `code_mode_prewarm` feature to establish the code-mode host session during startup, before the first turn.
- Make in-progress host initialization cancellable so a stalled prewarm does not block shutdown.
- Share a failed connection attempt with concurrent callers instead of immediately starting another attempt.

## Testing

- Verify app-server contacts the configured host before the first turn and can shut down while that connection is stalled.
- Exercise shared remote-host behavior with prewarming enabled.

GitOrigin-RevId: 3cfde5509be7b4aba80112c36c5616b9a9b632c6

codex-rs/app-server/tests/suite/v2/code_mode_host.rs

@@ -34,6 +34,35 @@ async fn app_server_shares_flag_selected_grpc_code_mode_host_across_threads() ->
     assert_shared_remote_code_mode_host("grpc://127.0.0.1:0").await
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn app_server_prewarms_flag_selected_code_mode_host_before_first_turn() -> Result<()> {
+    let model_server = responses::start_mock_server().await;
+    let codex_home = TempDir::new()?;
+    MockResponsesConfig::new(&model_server.uri())
+        .enable_feature(Feature::CodeModePrewarm)
+        .write(codex_home.path())?;
+
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+    let host_url = format!("ws://{}", listener.local_addr()?);
+    let mut app_server = TestAppServer::builder()
+        .with_codex_home(codex_home.path())
+        .with_args(&["--code-mode-host", &host_url])
+        .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT)
+        .await?;
+    app_server
+        .start_thread(ThreadStartParams::default())
+        .await?;
+
+    let (_stalled_connection, _) = timeout(DEFAULT_READ_TIMEOUT, listener.accept())
+        .await
+        .context("code-mode host was not contacted before the first turn")??;
+    let status = timeout(Duration::from_secs(5), app_server.shutdown_gracefully())
+        .await
+        .context("stalled code-mode prewarm blocked thread shutdown")??;
+    assert!(status.success(), "app-server did not exit successfully");
+    Ok(())
+}
+
 async fn assert_shared_remote_code_mode_host(listen_url: &str) -> Result<()> {
     let host_program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?;
     let mut code_mode_host = Command::new(host_program)
@@ -91,6 +120,7 @@ async fn assert_shared_remote_code_mode_host(listen_url: &str) -> Result<()> {
     let codex_home = TempDir::new()?;
     MockResponsesConfig::new(&model_server.uri())
         .enable_feature(Feature::CodeModeOnly)
+        .enable_feature(Feature::CodeModePrewarm)
         .write(codex_home.path())?;
     let original_config = std::fs::read_to_string(codex_home.path().join("config.toml"))?;
     let mut app_server = TestAppServer::builder()

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

@@ -184,6 +184,8 @@ struct OwnedCodeModeHost {
     endpoint: HostEndpoint,
     connection: StdMutex<Option<Arc<Connection>>>,
     connect_permit: Semaphore,
+    connection_generation: AtomicU64,
+    last_connection_error: StdMutex<Option<(u64, String)>>,
     next_session_id: AtomicU64,
 }
 
@@ -193,6 +195,8 @@ impl OwnedCodeModeHost {
             endpoint: HostEndpoint::Process(host_program),
             connection: StdMutex::new(None),
             connect_permit: Semaphore::new(/*permits*/ 1),
+            connection_generation: AtomicU64::new(0),
+            last_connection_error: StdMutex::new(None),
             next_session_id: AtomicU64::new(1),
         }
     }
@@ -205,6 +209,8 @@ impl OwnedCodeModeHost {
             },
             connection: StdMutex::new(None),
             connect_permit: Semaphore::new(/*permits*/ 1),
+            connection_generation: AtomicU64::new(0),
+            last_connection_error: StdMutex::new(None),
             next_session_id: AtomicU64::new(1),
         }
     }
@@ -214,18 +220,42 @@ impl OwnedCodeModeHost {
             return Ok(connection);
         }
 
+        let observed_generation = self.connection_generation.load(Ordering::Acquire);
         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 = match &self.endpoint {
-            HostEndpoint::Process(host_program) => Connection::spawn(host_program).await?,
+        let completed_generation = self.connection_generation.load(Ordering::Acquire);
+        if completed_generation != observed_generation
+            && let Some((generation, error)) = self
+                .last_connection_error
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .as_ref()
+            && *generation == completed_generation
+        {
+            return Err(ConnectionError::Other(error.clone()));
+        }
+        let 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?,
+            } => Connection::connect_websocket(websocket_url, http_client_factory).await,
+        };
+        let new_connection = match connection {
+            Ok(connection) => connection,
+            Err(error) => {
+                let generation = self.connection_generation.fetch_add(1, Ordering::AcqRel) + 1;
+                *self
+                    .last_connection_error
+                    .lock()
+                    .unwrap_or_else(std::sync::PoisonError::into_inner) =
+                    Some((generation, error.to_string()));
+                return Err(error);
+            }
         };
         let new_connection = Arc::new(new_connection);
         *self

codex-rs/core/config.schema.json

@@ -628,6 +628,9 @@
             "code_mode_only": {
               "type": "boolean"
             },
+            "code_mode_prewarm": {
+              "type": "boolean"
+            },
             "codex_git_commit": {
               "type": "boolean"
             },
@@ -5640,6 +5643,9 @@
         "code_mode_only": {
           "type": "boolean"
         },
+        "code_mode_prewarm": {
+          "type": "boolean"
+        },
         "codex_git_commit": {
           "type": "boolean"
         },

codex-rs/core/src/session_startup_prewarm.rs

@@ -17,6 +17,7 @@ use crate::responses_metadata::CodexResponsesRequestKind;
 use crate::session::INITIAL_SUBMIT_ID;
 use crate::session::session::Session;
 use crate::session::turn::build_prompt;
+use codex_features::Feature;
 use codex_otel::STARTUP_PREWARM_AGE_AT_FIRST_TURN_METRIC;
 use codex_otel::STARTUP_PREWARM_DURATION_METRIC;
 use codex_otel::SessionTelemetry;
@@ -183,6 +184,17 @@ impl SessionStartupPrewarmHandle {
 
 impl Session {
     pub(crate) async fn schedule_startup_prewarm(self: &Arc<Self>, base_instructions: String) {
+        if self.features().enabled(Feature::CodeModePrewarm)
+            && self.services.code_mode_service.is_available()
+        {
+            let session = Arc::clone(self);
+            tokio::spawn(async move {
+                if session.services.code_mode_service.session().await.is_err() {
+                    warn!("code-mode host startup prewarm failed");
+                }
+            });
+        }
+
         if !self.services.model_client.responses_websocket_enabled() {
             // Without websocket prewarm, resolve auth once so Agent Identity bootstrap can
             // register or engage this session's bearer fallback before the first user request.

codex-rs/core/src/tools/code_mode/mod.rs

@@ -72,7 +72,7 @@ pub(crate) struct CodeModeService {
     availability: Result<(), String>,
     dispatch_broker: Arc<CodeModeDispatchBroker>,
     default_exec_yield_time_ms: u64,
-    shutting_down: AtomicBool,
+    shutdown_token: CancellationToken,
     unavailable_warning_emitted: AtomicBool,
 }
 
@@ -89,7 +89,7 @@ impl CodeModeService {
             availability,
             dispatch_broker,
             default_exec_yield_time_ms: config.default_exec_yield_time_ms,
-            shutting_down: AtomicBool::new(false),
+            shutdown_token: CancellationToken::new(),
             unavailable_warning_emitted: AtomicBool::new(false),
         }
     }
@@ -160,7 +160,7 @@ impl CodeModeService {
     }
 
     pub(crate) async fn shutdown(&self) -> Result<(), String> {
-        self.shutting_down.store(true, Ordering::Release);
+        self.shutdown_token.cancel();
         // Join any initialization already in progress without initializing an unused service.
         match self
             .session
@@ -206,20 +206,25 @@ impl CodeModeService {
         )
     }
 
-    async fn session(&self) -> Result<Arc<dyn CodeModeSession>, String> {
-        if self.shutting_down.load(Ordering::Acquire) {
+    pub(crate) async fn session(&self) -> Result<Arc<dyn CodeModeSession>, String> {
+        if self.shutdown_token.is_cancelled() {
             return Err("code mode session is shutting down".to_string());
         }
         self.session
             .get_or_try_init(|| async {
-                if self.shutting_down.load(Ordering::Acquire) {
+                if self.shutdown_token.is_cancelled() {
                     return Err("code mode session is shutting down".to_string());
                 }
-                let session = self
-                    .session_provider
-                    .create_session(self.dispatch_broker.clone())
-                    .await?;
-                if self.shutting_down.load(Ordering::Acquire) {
+                let session = tokio::select! {
+                    biased;
+                    _ = self.shutdown_token.cancelled() => {
+                        return Err("code mode session is shutting down".to_string());
+                    }
+                    session = self
+                        .session_provider
+                        .create_session(self.dispatch_broker.clone()) => session?,
+                };
+                if self.shutdown_token.is_cancelled() {
                     let _ = session.shutdown().await;
                     return Err("code mode session is shutting down".to_string());
                 }

codex-rs/features/src/lib.rs

@@ -111,6 +111,8 @@ pub enum Feature {
     CodeModeBufferedExec,
     /// Run JavaScript code mode in the standalone host process.
     CodeModeHost,
+    /// Establish the code-mode host connection during session startup.
+    CodeModePrewarm,
     /// Terminate active code mode cells when their turn is interrupted.
     CodeModeInterrupt,
     /// Restrict model-visible tools to code mode entrypoints (`exec`, `wait`).
@@ -949,6 +951,12 @@ pub const FEATURES: &[FeatureSpec] = &[
         stage: Stage::Stable,
         default_enabled: true,
     },
+    FeatureSpec {
+        id: Feature::CodeModePrewarm,
+        key: "code_mode_prewarm",
+        stage: Stage::UnderDevelopment,
+        default_enabled: false,
+    },
     FeatureSpec {
         id: Feature::CodeModeInterrupt,
         key: "code_mode_interrupt",