From fb38afffa684b78f8e91f58f1f9d761789dc083f Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Fri, 12 Jun 2026 02:01:58 -0700 Subject: [PATCH] [codex] start threads with pending environments --- .../src/protocol/common.rs | 15 +- .../src/protocol/v2/environment.rs | 3 +- codex-rs/app-server/README.md | 2 +- .../environment_processor.rs | 13 +- .../app-server/tests/suite/v2/thread_start.rs | 92 +++++++++++ .../core/src/context/environment_context.rs | 26 ++- codex-rs/core/src/environment_selection.rs | 17 +- codex-rs/core/src/session/mod.rs | 5 +- codex-rs/core/src/session/session.rs | 2 - codex-rs/core/src/session/turn.rs | 16 +- codex-rs/core/src/session/turn_context.rs | 2 +- codex-rs/core/src/thread_manager.rs | 3 +- codex-rs/exec-server/src/client.rs | 154 +++++++++++++++++- codex-rs/exec-server/src/environment.rs | 72 +++++++- codex-rs/ext/skills/src/provider/executor.rs | 6 + 15 files changed, 380 insertions(+), 48 deletions(-) diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 0c2c88c747..b0e8c1d890 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -1955,7 +1955,7 @@ mod tests { request_id: request_id(), params: v2::EnvironmentAddParams { environment_id: "remote-a".to_string(), - exec_server_url: "ws://127.0.0.1:8765".to_string(), + exec_server_url: Some("ws://127.0.0.1:8765".to_string()), }, }; assert_eq!( @@ -2114,6 +2114,15 @@ mod tests { }), serde_json::to_value(&request)?, ); + assert_eq!( + serde_json::from_value::(json!({ + "environmentId": "pending-a" + }))?, + v2::EnvironmentAddParams { + environment_id: "pending-a".to_string(), + exec_server_url: None, + } + ); Ok(()) } @@ -2831,7 +2840,7 @@ mod tests { request_id: RequestId::Integer(9), params: v2::EnvironmentAddParams { environment_id: "remote-a".to_string(), - exec_server_url: "ws://127.0.0.1:8765".to_string(), + exec_server_url: Some("ws://127.0.0.1:8765".to_string()), }, }; assert_eq!( @@ -3212,7 +3221,7 @@ mod tests { request_id: RequestId::Integer(1), params: v2::EnvironmentAddParams { environment_id: "remote-a".to_string(), - exec_server_url: "ws://127.0.0.1:8765".to_string(), + exec_server_url: Some("ws://127.0.0.1:8765".to_string()), }, }; let reason = crate::experimental_api::ExperimentalApi::experimental_reason(&request); diff --git a/codex-rs/app-server-protocol/src/protocol/v2/environment.rs b/codex-rs/app-server-protocol/src/protocol/v2/environment.rs index 294ae736fd..d664c8e79d 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/environment.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/environment.rs @@ -8,7 +8,8 @@ use ts_rs::TS; #[ts(export_to = "v2/")] pub struct EnvironmentAddParams { pub environment_id: String, - pub exec_server_url: String, + #[ts(optional = nullable)] + pub exec_server_url: Option, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 1a820f6ffb..3f3964def6 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -196,7 +196,7 @@ Example with notification opt-out: - `experimentalFeature/list` — list feature flags with stage metadata (`beta`, `underDevelopment`, `stable`, etc.), enabled/default-enabled state, and cursor pagination. Pass `threadId` when showing feature state for an existing loaded thread so `enabled` is computed from that thread's refreshed config, including project-local config for the thread's cwd; if omitted, the server uses its default config resolution context. For non-beta flags, `displayName`/`description`/`announcement` are `null`. - `permissionProfile/list` — beta; list available permission profile ids with optional display `description` text, using cursor pagination. Pass `cwd` when the caller needs project-local `[permissions.]` entries to be included in the current catalog view. - `experimentalFeature/enablement/set` — patch the in-memory process-wide runtime feature enablement for currently supported feature keys. For each feature, precedence is: cloud requirements > --enable > config.toml > experimentalFeature/enablement/set (new) > code default. Invalid keys will be ignored. -- `environment/add` — experimental; add or replace a named remote environment by `environmentId` and `execServerUrl` for later selection by `thread/start` or `turn/start`; re-adding an `environmentId` updates existing selections, so use a new id for a different environment. Returns `{}` and does not change the default environment. +- `environment/add` — experimental; register a named remote environment by `environmentId` for later selection by `thread/start` or `turn/start`; omit `execServerUrl` while it is still starting, then call again with the URL when available. Re-adding an `environmentId` updates existing selections, so use a new id for a different environment. Returns `{}` and does not change the default environment. - `collaborationMode/list` — list available collaboration mode presets (experimental, no pagination). Built-in presets do not select a model; the Plan preset selects medium reasoning effort. This response omits built-in developer instructions; clients should either pass `settings.developer_instructions: null` when setting a mode to use Codex's built-in instructions, or provide their own instructions explicitly. - `skills/list` — list skills for one or more `cwd` values (optional `forceReload`). - `skills/extraRoots/set` — replace the app-server process runtime extra standalone skill roots. The roots are not persisted; missing directories are accepted and simply load no skills. diff --git a/codex-rs/app-server/src/request_processors/environment_processor.rs b/codex-rs/app-server/src/request_processors/environment_processor.rs index eb9b283f7b..0938683e5c 100644 --- a/codex-rs/app-server/src/request_processors/environment_processor.rs +++ b/codex-rs/app-server/src/request_processors/environment_processor.rs @@ -16,9 +16,16 @@ impl EnvironmentRequestProcessor { &self, params: EnvironmentAddParams, ) -> Result, JSONRPCErrorError> { - self.environment_manager - .upsert_environment(params.environment_id, params.exec_server_url) - .map_err(|err| invalid_request(err.to_string()))?; + match params.exec_server_url { + Some(exec_server_url) => self + .environment_manager + .upsert_environment(params.environment_id, exec_server_url), + None => self + .environment_manager + .register_pending_environment(params.environment_id) + .map(drop), + } + .map_err(|err| invalid_request(err.to_string()))?; Ok(Some(EnvironmentAddResponse {}.into())) } } diff --git a/codex-rs/app-server/tests/suite/v2/thread_start.rs b/codex-rs/app-server/tests/suite/v2/thread_start.rs index 82a2b6b2c0..c84ae34b84 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_start.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_start.rs @@ -7,6 +7,7 @@ use app_test_support::create_mock_responses_server_repeating_assistant; use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use codex_app_server_protocol::AskForApproval; +use codex_app_server_protocol::CapabilityRootLocation; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCResponse; @@ -14,6 +15,7 @@ use codex_app_server_protocol::McpServerStartupState; use codex_app_server_protocol::McpServerStatusUpdatedNotification; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::SandboxMode; +use codex_app_server_protocol::SelectedCapabilityRoot; use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ThreadSource; use codex_app_server_protocol::ThreadStartParams; @@ -312,6 +314,96 @@ async fn thread_start_rejects_unknown_environment_as_invalid_request() -> Result Ok(()) } +#[tokio::test] +async fn thread_start_and_first_turn_do_not_wait_for_pending_environment() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + let codex_home = TempDir::new()?; + create_config_toml_without_approval_policy(codex_home.path(), &server.uri())?; + + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let environment_request_id = mcp + .send_raw_request( + "environment/add", + Some(json!({"environmentId": "pending-environment"})), + ) + .await?; + let environment_response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(environment_request_id)), + ) + .await??; + assert_eq!(environment_response.result, json!({})); + + let environment = TurnEnvironmentParams { + environment_id: "pending-environment".to_string(), + cwd: codex_home.path().to_path_buf().try_into()?, + }; + let environment_cwd = std::fs::canonicalize(environment.cwd.as_path())?; + let thread_request_id = mcp + .send_thread_start_request(ThreadStartParams { + environments: Some(vec![environment]), + selected_capability_roots: Some(vec![SelectedCapabilityRoot { + id: "pending-skills".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "pending-environment".to_string(), + path: environment_cwd.to_string_lossy().into_owned(), + }, + }]), + ..Default::default() + }) + .await?; + let thread_response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_request_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response(thread_response)?; + + let turn_request_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id, + input: vec![V2UserInput::Text { + text: "start before the environment is ready".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_request_id)), + ) + .await??; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let requests = server + .received_requests() + .await + .context("failed to fetch received requests")?; + let model_request = requests + .iter() + .find(|request| request.url.path().ends_with("/responses")) + .context("expected model request")?; + let model_request_body = model_request + .body_json::() + .context("model request body should be JSON")? + .to_string(); + assert!(model_request_body.contains("still loading")); + assert!( + model_request_body.contains(&format!("{}", environment_cwd.display())), + "expected cwd {} in model request", + environment_cwd.display() + ); + + Ok(()) +} + #[tokio::test] async fn thread_start_response_includes_loaded_instruction_sources() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; diff --git a/codex-rs/core/src/context/environment_context.rs b/codex-rs/core/src/context/environment_context.rs index dd11ebd392..2e99cac7bd 100644 --- a/codex-rs/core/src/context/environment_context.rs +++ b/codex-rs/core/src/context/environment_context.rs @@ -15,6 +15,8 @@ use std::path::PathBuf; use super::ContextualUserFragment; +const ENVIRONMENT_SHELL_LOADING: &str = "still loading"; + #[derive(Debug, Clone, PartialEq)] pub(crate) struct EnvironmentContext { pub(crate) environments: EnvironmentContextEnvironments, @@ -47,11 +49,25 @@ impl EnvironmentContextEnvironment { .map(|environment| Self { id: environment.environment_id.clone(), cwd: environment.cwd.clone(), - shell: environment - .shell - .as_ref() - .map(|shell| shell.name().to_string()) - .unwrap_or_else(|| shell.name().to_string()), + shell: if environment.environment.is_remote() { + environment + .environment + .current_info() + .map(|info| info.shell.name) + .or_else(|| { + environment + .shell + .as_ref() + .map(|shell| shell.name().to_string()) + }) + .unwrap_or_else(|| ENVIRONMENT_SHELL_LOADING.to_string()) + } else { + environment + .shell + .as_ref() + .map(|shell| shell.name().to_string()) + .unwrap_or_else(|| shell.name().to_string()) + }, }) .collect() } diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 3ac1e43a3b..f2cf72b213 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -42,6 +42,7 @@ impl ResolvedTurnEnvironments { self.turn_environments.first() } + #[cfg(test)] pub(crate) fn primary_environment(&self) -> Option> { self.primary() .map(|environment| Arc::clone(&environment.environment)) @@ -49,6 +50,7 @@ impl ResolvedTurnEnvironments { pub(crate) fn primary_filesystem(&self) -> Option> { self.primary() + .filter(|environment| environment.environment.is_ready()) .map(|environment| environment.environment.get_filesystem()) } @@ -61,7 +63,7 @@ impl ResolvedTurnEnvironments { } } -pub(crate) async fn resolve_environment_selections( +pub(crate) fn resolve_environment_selections( environment_manager: &EnvironmentManager, environments: &[TurnEnvironmentSelection], ) -> CodexResult { @@ -80,8 +82,9 @@ pub(crate) async fn resolve_environment_selections( .ok_or_else(|| { CodexErr::InvalidRequest(format!("unknown turn environment id `{environment_id}`")) })?; - let shell = match environment.info().await { - Ok(info) => match Shell::from_environment_shell_info(info.shell) { + environment.start_loading_info(); + let shell = match environment.current_info() { + Some(info) => match Shell::from_environment_shell_info(info.shell) { Ok(shell) => Some(shell), Err(err) => { tracing::warn!( @@ -90,10 +93,7 @@ pub(crate) async fn resolve_environment_selections( None } }, - Err(err) => { - tracing::warn!("failed to get info for environment `{environment_id}`: {err}"); - None - } + None => None, }; turn_environments.push(TurnEnvironment { environment_id, @@ -205,7 +205,6 @@ url = "ws://127.0.0.1:8765" }, ], ) - .await .expect_err("duplicate environment id should fail"); assert!(err.to_string().contains("duplicate")); @@ -224,7 +223,6 @@ url = "ws://127.0.0.1:8765" cwd: selected_cwd, }], ) - .await .expect("environment selections should resolve"); assert_eq!( @@ -262,7 +260,6 @@ url = "ws://127.0.0.1:8765" cwd: cwd.clone(), }], ) - .await .expect("local environment should resolve"); let remote_environment = Arc::new( Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 08c691b9c1..f7df13d179 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -528,10 +528,7 @@ impl Codex { .startup_warnings .extend(user_instruction_provider_warnings); // TODO(anp) assemble instructions from multiple environments - let primary_environment = environment_selections.primary_environment(); - let primary_fs = primary_environment - .as_ref() - .map(|environment| environment.get_filesystem()); + let primary_fs = environment_selections.primary_filesystem(); let loaded_agents_md = load_project_instructions(&mut config, user_instructions, primary_fs.as_deref()).await; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index df60870a6e..b8d6d7a1bc 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -444,7 +444,6 @@ async fn warm_plugins_and_skills_for_session_init( environment_manager.as_ref(), &environments, ) - .await .ok() .and_then(|resolved| resolved.primary_filesystem()); let plugins_input = config.plugins_config_input(); @@ -1118,7 +1117,6 @@ impl Session { sess.services.environment_manager.as_ref(), session_configuration.environment_selections(), ) - .await .map_err(|err| { CodexErr::InvalidRequest(err.to_string().replace( "unknown turn environment id", diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index da1d825310..104ebbd30c 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -414,12 +414,16 @@ pub(crate) async fn run_turn( async fn turn_diff_display_roots(turn_context: &TurnContext) -> Vec<(String, PathBuf)> { let mut display_roots = Vec::new(); for turn_environment in &turn_context.environments.turn_environments { - let root = get_git_repo_root_with_fs( - turn_environment.environment.get_filesystem().as_ref(), - &turn_environment.cwd, - ) - .await - .unwrap_or_else(|| turn_environment.cwd.clone()) + let root = if turn_environment.environment.is_ready() { + get_git_repo_root_with_fs( + turn_environment.environment.get_filesystem().as_ref(), + &turn_environment.cwd, + ) + .await + .unwrap_or_else(|| turn_environment.cwd.clone()) + } else { + turn_environment.cwd.clone() + } .into_path_buf(); display_roots.push((turn_environment.environment_id.clone(), root)); } diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index c249d713ab..9d5ee7babe 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -704,7 +704,6 @@ impl Session { self.services.environment_manager.as_ref(), session_configuration.environment_selections(), ) - .await .unwrap_or_else(|err| { warn!("failed to resolve turn environments: {err}"); ResolvedTurnEnvironments::default() @@ -748,6 +747,7 @@ impl Session { let effective_skill_roots = plugin_outcome.effective_plugin_skill_roots(); let skills_input = skills_load_input_from_config(&per_turn_config, effective_skill_roots); let fs = primary_turn_environment + .filter(|turn_environment| turn_environment.environment.is_ready()) .map(|turn_environment| turn_environment.environment.get_filesystem()); let skills_outcome = Arc::new( self.services diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index c332622acb..2add5ded5f 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1382,8 +1382,7 @@ impl ThreadManagerState { } } let environment_selections = - resolve_environment_selections(self.environment_manager.as_ref(), &environments) - .await?; + resolve_environment_selections(self.environment_manager.as_ref(), &environments)?; let user_instructions = self .user_instructions_for_spawn(&session_source, parent_thread_id, forked_from_thread_id) .await; diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index e9a03875f7..39a8e52a70 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use std::sync::Mutex as StdMutex; use std::sync::OnceLock; use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use std::time::Duration; use arc_swap::ArcSwap; @@ -84,6 +85,9 @@ use crate::protocol::TerminateParams; use crate::protocol::TerminateResponse; use crate::protocol::WriteParams; use crate::protocol::WriteResponse; + +const ENVIRONMENT_INFO_RELOAD_PENDING: u64 = 1 << 63; +const ENVIRONMENT_INFO_LOAD_TOKEN_MASK: u64 = !ENVIRONMENT_INFO_RELOAD_PENDING; use crate::rpc::RpcCallError; use crate::rpc::RpcClient; use crate::rpc::RpcClientEvent; @@ -203,12 +207,15 @@ pub(crate) struct LazyRemoteExecServerClient { state: Arc>, transport_revision_tx: watch::Sender, connect_lock: Arc, + environment_info_lock: Arc, + environment_info_load_state: Arc, } struct LazyRemoteExecServerClientState { // None means we do not know how to connect to this environment yet. transport_params: Option>, client: Option, + environment_info: Option, } impl LazyRemoteExecServerClient { @@ -227,9 +234,12 @@ impl LazyRemoteExecServerClient { state: Arc::new(StdMutex::new(LazyRemoteExecServerClientState { transport_params: transport_params.map(Arc::new), client: None, + environment_info: None, })), transport_revision_tx, connect_lock: Arc::new(Semaphore::new(/*permits*/ 1)), + environment_info_lock: Arc::new(Semaphore::new(/*permits*/ 1)), + environment_info_load_state: Arc::new(AtomicU64::new(0)), } } @@ -306,6 +316,7 @@ impl LazyRemoteExecServerClient { } state.transport_params = Some(Arc::new(transport_params)); state.client = None; + state.environment_info = None; drop(state); // Notify after releasing the state lock so woken callers can inspect it immediately. self.transport_revision_tx @@ -380,7 +391,116 @@ impl HttpClient for LazyRemoteExecServerClient { impl LazyRemoteExecServerClient { pub(crate) async fn environment_info(&self) -> Result { - self.get().await?.environment_info().await + if let Some(info) = self.current_environment_info() { + return Ok(info); + } + let _info_permit = + self.environment_info_lock.acquire().await.map_err(|_| { + ExecServerError::Protocol("environment info lock closed".to_string()) + })?; + loop { + if let Some(info) = self.current_environment_info() { + return Ok(info); + } + let transport_params = self.wait_for_transport().await?; + let client = self.get().await?; + let mut transport_revision_rx = self.transport_revision_tx.subscribe(); + if !self.is_current_transport(&transport_params) { + continue; + } + let info = tokio::select! { + result = client.environment_info() => { + match result { + Ok(info) => info, + Err(_err) if !self.is_current_transport(&transport_params) => continue, + Err(err) => return Err(err), + } + } + changed = transport_revision_rx.changed() => { + changed.map_err(|_| ExecServerError::Closed)?; + continue; + } + }; + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .transport_params + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, &transport_params)) + { + state.environment_info = Some(info.clone()); + return Ok(info); + } + } + } + + pub(crate) fn current_environment_info(&self) -> Option { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .environment_info + .clone() + } + + pub(crate) fn try_begin_environment_info_load(&self) -> Option { + if self.current_environment_info().is_some() { + return None; + } + let token = self.environment_info_load_token(); + let mut state = self.environment_info_load_state.load(Ordering::Acquire); + loop { + if state == 0 { + match self.environment_info_load_state.compare_exchange( + 0, + token, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return Some(token), + Err(next_state) => state = next_state, + } + } else { + let pending_state = state | ENVIRONMENT_INFO_RELOAD_PENDING; + if pending_state == state { + return None; + } + match self.environment_info_load_state.compare_exchange( + state, + pending_state, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return None, + Err(next_state) => state = next_state, + } + } + } + } + + pub(crate) fn finish_environment_info_load(&self, token: u64) -> bool { + let mut state = self.environment_info_load_state.load(Ordering::Acquire); + loop { + debug_assert_eq!(state & ENVIRONMENT_INFO_LOAD_TOKEN_MASK, token); + let reload_requested = state & ENVIRONMENT_INFO_RELOAD_PENDING != 0; + match self.environment_info_load_state.compare_exchange( + state, + 0, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return reload_requested, + Err(next_state) => state = next_state, + } + } + } + + pub(crate) fn environment_info_load_token(&self) -> u64 { + self.transport_revision_tx + .borrow() + .wrapping_add(1) + .min(ENVIRONMENT_INFO_LOAD_TOKEN_MASK) } } @@ -1074,6 +1194,7 @@ mod tests { use tokio::io::duplex; use tokio::net::TcpListener; use tokio::net::TcpStream; + use tokio::sync::Barrier; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::time::Duration; @@ -1759,6 +1880,37 @@ mod tests { ready_server.await.expect("ready server should finish"); } + #[tokio::test] + async fn environment_info_load_preserves_concurrent_reload_handoff() { + for _ in 0..100 { + let client = Arc::new(LazyRemoteExecServerClient::new( + ExecServerTransportParams::WebSocketUrl { + websocket_url: "ws://127.0.0.1:8765".to_string(), + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + }, + )); + let first_token = client + .try_begin_environment_info_load() + .expect("first load should start"); + let barrier = Arc::new(Barrier::new(2)); + let contender = { + let barrier = Arc::clone(&barrier); + let client = Arc::clone(&client); + tokio::spawn(async move { + barrier.wait().await; + client.try_begin_environment_info_load() + }) + }; + + barrier.wait().await; + let reload_requested = client.finish_environment_info_load(first_token); + let contender_token = contender.await.expect("contender should not panic"); + + assert!(reload_requested || contender_token.is_some()); + } + } + #[tokio::test] async fn wake_notifications_do_not_block_other_sessions() { let (client_stdin, server_reader) = duplex(1 << 20); diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 9ecb01733e..16eba9811d 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -296,16 +296,19 @@ impl EnvironmentManager { .unwrap_or_else(std::sync::PoisonError::into_inner); match environments.get(&environment_id) { // Update the shared client inside the stable Environment handle. - Some(existing) if existing.is_remote() => existing - .set_remote_transport(ExecServerTransportParams::websocket_url(exec_server_url)), + Some(existing) if existing.is_remote() => { + existing.set_remote_transport(ExecServerTransportParams::websocket_url( + exec_server_url, + )); + existing.start_loading_info(); + } _ => { - environments.insert( - environment_id, - Arc::new(Environment::remote_inner( - exec_server_url, - self.local_runtime_paths.clone(), - )), - ); + let environment = Arc::new(Environment::remote_inner( + exec_server_url, + self.local_runtime_paths.clone(), + )); + environment.start_loading_info(); + environments.insert(environment_id, environment); } } Ok(()) @@ -519,6 +522,42 @@ impl Environment { self.info_provider.info().await } + /// Returns already-loaded environment information without starting a connection. + pub fn current_info(&self) -> Option { + self.remote_client.as_ref().map_or_else( + || Some(EnvironmentInfo::local()), + LazyRemoteExecServerClient::current_environment_info, + ) + } + + /// Starts loading remote metadata without making the caller wait for it. + pub fn start_loading_info(self: &Arc) { + if self.current_info().is_some() || !self.is_ready() { + return; + } + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + return; + }; + let Some(remote_client) = self.remote_client.clone() else { + return; + }; + let Some(load_token) = remote_client.try_begin_environment_info_load() else { + return; + }; + let environment = Arc::clone(self); + runtime.spawn(async move { + let result = environment.info().await; + let transport_changed = remote_client.environment_info_load_token() != load_token; + let reload_requested = remote_client.finish_environment_info_load(load_token); + if let Err(err) = result { + tracing::warn!("failed to load environment metadata: {err}"); + } + if transport_changed || reload_requested { + environment.start_loading_info(); + } + }); + } + pub fn get_exec_backend(&self) -> Arc { Arc::clone(&self.exec_backend) } @@ -964,6 +1003,20 @@ mod tests { first.exec_server_url().as_deref(), Some(exec_server_url.as_str()) ); + let loaded_info = tokio::time::timeout(Duration::from_secs(1), async { + loop { + if let Some(info) = first.current_info() { + break info; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("environment metadata should load in the background"); + assert_eq!( + first.info().await.expect("cached environment info"), + loaded_info + ); assert_eq!(manager.default_environment_id(), None); manager @@ -977,6 +1030,7 @@ mod tests { second.exec_server_url().as_deref(), Some("ws://127.0.0.1:9876") ); + assert_eq!(second.current_info(), None); assert!(Arc::ptr_eq(&first, &second)); assert!(Arc::ptr_eq(&exec_backend, &second.get_exec_backend())); assert!(Arc::ptr_eq(&http_client, &second.get_http_client())); diff --git a/codex-rs/ext/skills/src/provider/executor.rs b/codex-rs/ext/skills/src/provider/executor.rs index 1ed1d11fa0..36b1a41ac1 100644 --- a/codex-rs/ext/skills/src/provider/executor.rs +++ b/codex-rs/ext/skills/src/provider/executor.rs @@ -65,6 +65,12 @@ impl SkillProvider for ExecutorSkillProvider { )); continue; }; + if !environment.is_ready() { + catalog.warnings.push(format!( + "Selected capability root `{selected_root_id}` is unavailable while environment `{environment_id}` is starting." + )); + continue; + } let root_path = match executor_absolute_path(&path) { Ok(root_path) => root_path, Err(err) => {