diff --git a/codex-rs/app-server/tests/suite/v2/code_mode_host.rs b/codex-rs/app-server/tests/suite/v2/code_mode_host.rs index 534c34878f..82d2a7fa50 100644 --- a/codex-rs/app-server/tests/suite/v2/code_mode_host.rs +++ b/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() diff --git a/codex-rs/code-mode/src/remote_session.rs b/codex-rs/code-mode/src/remote_session.rs index 502642129f..5433b88f38 100644 --- a/codex-rs/code-mode/src/remote_session.rs +++ b/codex-rs/code-mode/src/remote_session.rs @@ -184,6 +184,8 @@ struct OwnedCodeModeHost { endpoint: HostEndpoint, connection: StdMutex>>, connect_permit: Semaphore, + connection_generation: AtomicU64, + last_connection_error: StdMutex>, 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 diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 7f6d6c2666..15d517d398 100644 --- a/codex-rs/core/config.schema.json +++ b/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" }, diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index 4ae1af8b85..70a9efea2f 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/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, 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. diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 94e9ce7c1b..c6a183a31a 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -72,7 +72,7 @@ pub(crate) struct CodeModeService { availability: Result<(), String>, dispatch_broker: Arc, 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, String> { - if self.shutting_down.load(Ordering::Acquire) { + pub(crate) async fn session(&self) -> Result, 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()); } diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index fd00597734..26548db324 100644 --- a/codex-rs/features/src/lib.rs +++ b/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",