openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Resolve selected roots through current environments

Browse Source
This commit is contained in:
jif-oai
2026-06-25 14:03:28 +01:00
parent cfa422c672
commit 740c8d8fcc
5 changed files with 183 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/environment_selection.rs
+13 -69
View File
@@ -1,3 +1,4 @@
use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt;
use std::sync::Arc;
@@ -7,9 +8,6 @@ use codex_exec_server::Environment;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::ExecServerError;
use codex_exec_server::ExecutorFileSystem;
use codex_exec_server::ResolvedSelectedCapabilityRoot;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_path_uri::PathUri;
@@ -230,24 +228,20 @@ pub(crate) struct TurnEnvironmentSnapshot {
}
impl TurnEnvironmentSnapshot {
pub(crate) fn resolve_selected_capability_roots(
&self,
selected_roots: &[SelectedCapabilityRoot],
) -> Vec<ResolvedSelectedCapabilityRoot> {
selected_roots
pub(crate) fn captured_environments(&self) -> HashMap<String, Option<Arc<Environment>>> {
self.turn_environments
.iter()
.filter_map(|selected_root| {
let CapabilityRootLocation::Environment { environment_id, .. } =
&selected_root.location;
let environment = self
.turn_environments
.iter()
.find(|environment| environment.environment_id == *environment_id)?;
Some(ResolvedSelectedCapabilityRoot::new(
selected_root.clone(),
Arc::clone(&environment.environment),
))
.map(|environment| {
(
environment.environment_id.clone(),
Some(Arc::clone(&environment.environment)),
)
})
.chain(
self.starting
.iter()
.map(|environment| (environment.selection.environment_id.clone(), None)),
)
.collect()
}
@@ -584,56 +578,6 @@ url = "ws://127.0.0.1:8765"
assert_eq!(resolved.snapshot().await.to_selections(), vec![local]);
}
#[tokio::test]
async fn selected_capability_roots_bind_only_to_their_ready_environment() {
let cwd = PathUri::parse("file:///workspace").expect("cwd URI");
let primary = Arc::new(Environment::default_for_tests());
let selected = Arc::new(Environment::default_for_tests());
let snapshot = TurnEnvironmentSnapshot {
turn_environments: vec![
TurnEnvironment::new(
"primary".to_string(),
Arc::clone(&primary),
cwd.clone(),
/*shell*/ None,
),
TurnEnvironment::new(
"selected".to_string(),
Arc::clone(&selected),
cwd,
/*shell*/ None,
),
],
starting: Vec::new(),
};
let plugin_path = PathUri::parse("file:///plugins/demo").expect("plugin URI");
let roots = [
SelectedCapabilityRoot {
id: "demo@1".to_string(),
location: CapabilityRootLocation::Environment {
environment_id: "selected".to_string(),
path: plugin_path.clone(),
},
},
SelectedCapabilityRoot {
id: "pending@1".to_string(),
location: CapabilityRootLocation::Environment {
environment_id: "pending".to_string(),
path: plugin_path,
},
},
];
let resolved = snapshot.resolve_selected_capability_roots(&roots);
let [resolved] = resolved.as_slice() else {
panic!("only the ready selected environment should resolve");
};
assert_eq!(resolved.selected_root(), &roots[0]);
assert!(Arc::ptr_eq(resolved.environment(), &selected));
assert!(!Arc::ptr_eq(resolved.environment(), &primary));
}
#[tokio::test]
async fn blocking_snapshot_waits_for_starting_environment() {
let listener = TcpListener::bind("127.0.0.1:0")
codex-rs/core/src/session/mod.rs
+9 -2
View File
@@ -2828,8 +2828,15 @@ impl Session {
.await;
}
let loaded_agents_md = self.services.agents_md_manager.get_loaded().await;
let selected_capability_roots = environments
.resolve_selected_capability_roots(&self.services.selected_capability_roots);
let selected_capability_roots = self
.services
.turn_environments
.environment_manager()
.resolve_selected_capability_roots(
&self.services.selected_capability_roots,
&environments.captured_environments(),
)
.await;
Arc::new(StepContext::new(
turn_context,
environments,
codex-rs/exec-server/src/environment.rs
+20 -1
View File
@@ -54,7 +54,7 @@ pub const CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID_ENV_VAR: &str =
#[derive(Debug)]
pub struct EnvironmentManager {
default_environment: Option<String>,
environments: RwLock<HashMap<String, Arc<Environment>>>,
pub(super) environments: RwLock<HashMap<String, Arc<Environment>>>,
local_environment: Option<Arc<Environment>>,
local_runtime_paths: Option<ExecServerRuntimePaths>,
}
@@ -574,6 +574,25 @@ impl Environment {
}
}
/// Starts the initial connection after an environment is actually selected for use.
pub(crate) fn start_connecting_for_use(environment: &Arc<Self>) {
if environment.remote_client.is_none() {
return;
}
let mut startup_task = environment
.startup_task
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if startup_task.is_none() {
let environment = Arc::clone(environment);
*startup_task = Some(AbortOnDropHandle::new(tokio::spawn(async move {
if let Err(error) = environment.wait_until_ready().await {
tracing::debug!(%error, "exec-server environment startup failed");
}
})));
}
}
/// Returns whether initial startup has either succeeded or permanently failed.
pub fn startup_finished(&self) -> bool {
self.remote_client
codex-rs/exec-server/src/resolved_capability.rs
+70 -1
View File
@@ -1,9 +1,12 @@
use std::collections::HashMap;
use std::fmt;
use std::sync::Arc;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use crate::Environment;
use crate::EnvironmentManager;
use crate::ExecutorFileSystem;
/// A selected capability root pinned to the exact environment instance that owns it.
@@ -17,7 +20,7 @@ pub struct ResolvedSelectedCapabilityRoot {
}
impl ResolvedSelectedCapabilityRoot {
pub fn new(selected_root: SelectedCapabilityRoot, environment: Arc<Environment>) -> Self {
fn new(selected_root: SelectedCapabilityRoot, environment: Arc<Environment>) -> Self {
Self {
selected_root,
environment,
@@ -37,6 +40,72 @@ impl ResolvedSelectedCapabilityRoot {
}
}
impl EnvironmentManager {
/// Binds selected roots to the ready environment instances currently registered for them.
///
/// The registry and readiness are each captured once per environment ID, so every root owned
/// by one environment uses the same executor instance and readiness result. Missing, starting,
/// or failed environments are omitted. A lazy environment is started for a later step.
pub async fn resolve_selected_capability_roots(
&self,
selected_roots: &[SelectedCapabilityRoot],
captured_environments: &HashMap<String, Option<Arc<Environment>>>,
) -> Vec<ResolvedSelectedCapabilityRoot> {
let candidates = {
let environments = self
.environments
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
selected_roots
.iter()
.filter_map(|selected_root| {
let CapabilityRootLocation::Environment { environment_id, .. } =
&selected_root.location;
let (environment, already_ready) =
match captured_environments.get(environment_id) {
Some(Some(environment)) => (Arc::clone(environment), true),
Some(None) => return None,
None => (Arc::clone(environments.get(environment_id)?), false),
};
Some((
ResolvedSelectedCapabilityRoot::new(selected_root.clone(), environment),
already_ready,
))
})
.collect::<Vec<_>>()
};
let mut readiness = HashMap::new();
for (candidate, already_ready) in &candidates {
let CapabilityRootLocation::Environment { environment_id, .. } =
&candidate.selected_root().location;
if readiness.contains_key(environment_id) {
continue;
}
let environment = candidate.environment();
let ready = if *already_ready {
true
} else if environment.startup_finished() {
environment.wait_until_ready().await.is_ok()
} else {
Environment::start_connecting_for_use(environment);
false
};
readiness.insert(environment_id.clone(), ready);
}
candidates
.into_iter()
.map(|(candidate, _)| candidate)
.filter(|candidate| {
let CapabilityRootLocation::Environment { environment_id, .. } =
&candidate.selected_root().location;
readiness.get(environment_id).copied().unwrap_or(false)
})
.collect()
}
}
impl fmt::Debug for ResolvedSelectedCapabilityRoot {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
codex-rs/exec-server/tests/selected_capability_roots.rs
+71
View File
@@ -0,0 +1,71 @@
#![cfg(unix)]
mod common;
use std::collections::HashMap;
use std::sync::Arc;
use codex_exec_server::EnvironmentManager;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_utils_path_uri::PathUri;
use common::exec_server::exec_server;
use pretty_assertions::assert_eq;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn selected_capability_roots_follow_same_id_environment_replacement() -> anyhow::Result<()> {
let mut executor_a = exec_server().await?;
let mut executor_b = exec_server().await?;
let manager = EnvironmentManager::without_environments();
let selected_root = SelectedCapabilityRoot {
id: "demo@1".to_string(),
location: CapabilityRootLocation::Environment {
environment_id: "tools".to_string(),
path: PathUri::parse("file:///plugins/demo")?,
},
};
manager.upsert_environment(
"tools".to_string(),
executor_a.websocket_url().to_string(),
/*connect_timeout*/ None,
)?;
let environment_a = manager
.get_environment("tools")
.expect("executor A should be registered");
environment_a.wait_until_ready().await?;
let resolved_a = manager
.resolve_selected_capability_roots(std::slice::from_ref(&selected_root), &HashMap::new())
.await;
let [resolved_a] = resolved_a.as_slice() else {
anyhow::bail!("selected root should resolve through executor A");
};
manager.upsert_environment(
"tools".to_string(),
executor_b.websocket_url().to_string(),
/*connect_timeout*/ None,
)?;
let environment_b = manager
.get_environment("tools")
.expect("executor B should be registered");
environment_b.wait_until_ready().await?;
let resolved_b = manager
.resolve_selected_capability_roots(std::slice::from_ref(&selected_root), &HashMap::new())
.await;
let [resolved_b] = resolved_b.as_slice() else {
anyhow::bail!("selected root should resolve through executor B");
};
assert_eq!(resolved_b.selected_root(), &selected_root);
assert!(Arc::ptr_eq(resolved_a.environment(), &environment_a));
assert!(Arc::ptr_eq(resolved_b.environment(), &environment_b));
assert!(!Arc::ptr_eq(
resolved_a.environment(),
resolved_b.environment()
));
executor_a.shutdown().await?;
executor_b.shutdown().await?;
Ok(())
}