mirror of
https://github.com/openai/codex.git
synced 2026-11-06 20:20:34 +00:00
[codex] start threads with pending environments
This commit is contained in:
@@ -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::<v2::EnvironmentAddParams>(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);
|
||||
|
||||
@@ -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<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
|
||||
|
||||
@@ -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.<id>]` 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 <feature_name> > 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.
|
||||
|
||||
@@ -16,9 +16,16 @@ impl EnvironmentRequestProcessor {
|
||||
&self,
|
||||
params: EnvironmentAddParams,
|
||||
) -> Result<Option<ClientResponsePayload>, 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()))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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::<Value>()
|
||||
.context("model request body should be JSON")?
|
||||
.to_string();
|
||||
assert!(model_request_body.contains("<shell>still loading</shell>"));
|
||||
assert!(
|
||||
model_request_body.contains(&format!("<cwd>{}</cwd>", 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;
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -42,6 +42,7 @@ impl ResolvedTurnEnvironments {
|
||||
self.turn_environments.first()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn primary_environment(&self) -> Option<Arc<codex_exec_server::Environment>> {
|
||||
self.primary()
|
||||
.map(|environment| Arc::clone(&environment.environment))
|
||||
@@ -49,6 +50,7 @@ impl ResolvedTurnEnvironments {
|
||||
|
||||
pub(crate) fn primary_filesystem(&self) -> Option<Arc<dyn ExecutorFileSystem>> {
|
||||
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<ResolvedTurnEnvironments> {
|
||||
@@ -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()))
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<StdMutex<LazyRemoteExecServerClientState>>,
|
||||
transport_revision_tx: watch::Sender<u64>,
|
||||
connect_lock: Arc<Semaphore>,
|
||||
environment_info_lock: Arc<Semaphore>,
|
||||
environment_info_load_state: Arc<AtomicU64>,
|
||||
}
|
||||
|
||||
struct LazyRemoteExecServerClientState {
|
||||
// None means we do not know how to connect to this environment yet.
|
||||
transport_params: Option<Arc<ExecServerTransportParams>>,
|
||||
client: Option<ExecServerClient>,
|
||||
environment_info: Option<EnvironmentInfo>,
|
||||
}
|
||||
|
||||
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<EnvironmentInfo, ExecServerError> {
|
||||
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<EnvironmentInfo> {
|
||||
self.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.environment_info
|
||||
.clone()
|
||||
}
|
||||
|
||||
pub(crate) fn try_begin_environment_info_load(&self) -> Option<u64> {
|
||||
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);
|
||||
|
||||
@@ -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<EnvironmentInfo> {
|
||||
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<Self>) {
|
||||
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<dyn ExecBackend> {
|
||||
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()));
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
Reference in New Issue
Block a user