openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-18 17:18:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refresh resumed thread capability roots from executors (#39746)

Browse Source 
## Why

Resumed threads can carry a persisted capability-root location that no longer
matches the location reported by the attached executor.

## What changed

- Refresh matching thread-owned capability roots from the live environment while
  preserving persisted roots when the executor reports none.
- Use the same merge behavior for capability-root inspection and per-step MCP
  resolution, without changing owner-configured root precedence.

## Testing

Extend the remote-environment resume test to cover both an executor-provided
replacement root and the persisted-root fallback.

GitOrigin-RevId: 0c66195b2406f431a2e156e00d5225a8d22d18b2
This commit is contained in:
sayan-oai
2026-08-20 17:12:15 +00:00
committed by copyberry
parent c3db180493
commit 2e1f18e0db
3 changed files with 167 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/environment_selection.rs
+67 -7
View File
@@ -54,6 +54,59 @@ impl EnvironmentConfigOrigin {
}
selection
}
pub(crate) fn selected_capability_roots(
self,
environment: &Environment,
config: &EnvironmentConfig,
) -> Vec<SelectedCapabilityRoot> {
match self {
Self::Thread => environment.selected_capability_roots(),
Self::Owner => config.selected_capability_roots.clone(),
}
}
}
/// Combines persisted thread roots with attachment roots in selection order.
///
/// Thread-owned attachments still rely on roots reported by legacy executors. A matching live root
/// refreshes the persisted location while explicit owner configuration retains the existing
/// thread-first collision behavior.
pub(crate) fn combine_selected_capability_roots(
thread_roots: &[SelectedCapabilityRoot],
sources: impl IntoIterator<Item = (EnvironmentConfigOrigin, Vec<SelectedCapabilityRoot>)>,
) -> Vec<SelectedCapabilityRoot> {
let attachment_roots = sources
.into_iter()
.flat_map(|(config_origin, roots)| roots.into_iter().map(move |root| (config_origin, root)))
.collect::<Vec<_>>();
let mut combined_roots = thread_roots
.iter()
.map(|thread_root| {
let CapabilityRootLocation::Environment {
environment_id: thread_environment_id,
..
} = &thread_root.location;
attachment_roots
.iter()
.find_map(|(origin, attachment_root)| {
if *origin != EnvironmentConfigOrigin::Thread
|| attachment_root.id != thread_root.id
{
return None;
}
let CapabilityRootLocation::Environment {
environment_id: attachment_environment_id,
..
} = &attachment_root.location;
(attachment_environment_id == thread_environment_id).then_some(attachment_root)
})
.unwrap_or(thread_root)
.clone()
})
.collect::<Vec<_>>();
combined_roots.extend(attachment_roots.into_iter().map(|(_, root)| root));
combined_roots
}
pub(crate) fn default_thread_environment_selections(
@@ -410,13 +463,20 @@ impl ThreadEnvironments {
thread_selected_capability_roots: &[SelectedCapabilityRoot],
) -> SelectedCapabilityRootsStatus {
let environments = self.environments.load();
let mut selected_capability_roots = thread_selected_capability_roots.to_vec();
for environment in environments.iter() {
let EnvironmentConfigState::Ready(config) = &environment.selection.config else {
continue;
};
selected_capability_roots.extend(config.selected_capability_roots.iter().cloned());
}
let mut selected_capability_roots = combine_selected_capability_roots(
thread_selected_capability_roots,
environments.iter().filter_map(|environment| {
let EnvironmentConfigState::Ready(config) = &environment.selection.config else {
return None;
};
Some((
environment.config_origin,
environment
.config_origin
.selected_capability_roots(&environment.environment, config),
))
}),
);
let mut seen_root_ids = HashSet::with_capacity(selected_capability_roots.len());
selected_capability_roots.retain(|root| seen_root_ids.insert(root.id.clone()));
codex-rs/core/src/session/mcp.rs
+13 -12
View File
@@ -1,5 +1,6 @@
use super::mcp_refresh::McpRefreshInvalidationGuard;
use super::*;
use crate::environment_selection::combine_selected_capability_roots;
use crate::tools::sandboxing::executor_windows_sandbox_level;
use codex_exec_server::ExecutorCapabilityDiscoveryCache;
use codex_exec_server::ExecutorCapabilityDiscoverySnapshot;
@@ -438,18 +439,18 @@ impl Session {
let mut root_locations_by_id = HashMap::new();
let mut selected_capability_roots = Vec::new();
let mut ready_environment_root_count = 0;
for (index, root) in self
.services
.selected_capability_roots
.iter()
.cloned()
.chain(
environments
.turn_environments()
.flat_map(|environment| environment.config().selected_capability_roots.clone()),
)
.enumerate()
{
let combined_roots = combine_selected_capability_roots(
&self.services.selected_capability_roots,
environments.turn_environments().map(|environment| {
(
environment.config_origin,
environment
.config_origin
.selected_capability_roots(&environment.environment, environment.config()),
)
}),
);
for (index, root) in combined_roots.into_iter().enumerate() {
if let Some(kept_location) = root_locations_by_id.get(&root.id) {
if kept_location != &root.location {
tracing::warn!(
codex-rs/core/tests/suite/remote_env.rs
+87 -28
View File
@@ -118,6 +118,7 @@ use std::fs;
use std::path::PathBuf;
use std::process::Command;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;
@@ -158,13 +159,22 @@ impl ThreadLifecycleContributor<Config> for WaitForEnvironmentTestExtension {
}
}
struct ReadyCapabilityRootsTestExtension;
#[derive(Default)]
struct ReadyCapabilityRootsTestExtension {
observed_roots: Option<Arc<Mutex<Vec<Vec<SelectedCapabilityRoot>>>>>,
}
impl ContextContributor for ReadyCapabilityRootsTestExtension {
fn contribute_world_state<'a>(
&'a self,
input: WorldStateContributionInput<'a>,
) -> ExtensionFuture<'a, Vec<WorldStateSectionContribution>> {
if let Some(observed_roots) = &self.observed_roots {
observed_roots
.lock()
.expect("observed capability roots should not be poisoned")
.push(input.ready_selected_capability_roots.to_vec());
}
let root_ids = input
.ready_selected_capability_roots
.iter()
@@ -1218,7 +1228,7 @@ async fn wait_for_response_request_count(response_mock: &ResponseMock, expected_
async fn shared_executor_keeps_ready_capability_roots_scoped_to_each_attachment() -> Result<()> {
let server = start_mock_server().await;
let mut extensions = ExtensionRegistryBuilder::new();
extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension));
extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension::default()));
let mut builder = test_codex()
.with_extensions(Arc::new(extensions.build()))
.with_config(|config| {
@@ -1484,7 +1494,7 @@ async fn pending_attachment_installs_configuration_before_waiting_turn_resumes()
let server = start_mock_server().await;
let mut extensions = ExtensionRegistryBuilder::new();
extensions.thread_lifecycle_contributor(Arc::new(WaitForEnvironmentTestExtension));
extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension));
extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension::default()));
let mut builder = test_codex()
.with_extensions(Arc::new(extensions.build()))
.with_config(|config| {
@@ -1718,9 +1728,12 @@ async fn pending_attachment_installs_configuration_before_waiting_turn_resumes()
Ok(())
}
#[test_case(true; "uses refreshed executor root")]
#[test_case(false; "preserves persisted root when executor reports none")]
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn ready_before_selection_exposes_remote_tools_and_capability_context_after_wait()
-> Result<()> {
async fn ready_before_selection_resolves_resumed_thread_capability_root_after_wait(
executor_reports_refreshed_root: bool,
) -> Result<()> {
const WAIT_CALL_ID: &str = "wait-ready-before-selection";
let rendezvous = TcpListener::bind("127.0.0.1:0").await?;
@@ -1804,9 +1817,12 @@ async fn ready_before_selection_exposes_remote_tools_and_capability_context_afte
],
)
.await;
let observed_roots = Arc::new(Mutex::new(Vec::new()));
let mut extensions = ExtensionRegistryBuilder::new();
extensions.thread_lifecycle_contributor(Arc::new(WaitForEnvironmentTestExtension));
extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension));
extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension {
observed_roots: Some(Arc::clone(&observed_roots)),
}));
let mut builder = test_codex()
.with_extensions(Arc::new(extensions.build()))
.with_config(|config| {
@@ -1816,20 +1832,36 @@ async fn ready_before_selection_exposes_remote_tools_and_capability_context_afte
assert!(config.features.enable(Feature::UnifiedExec).is_ok());
});
let test = builder.build(&server).await?;
let ready_root = SelectedCapabilityRoot {
let refreshed_root = SelectedCapabilityRoot {
id: "ready-first-root".to_string(),
location: CapabilityRootLocation::Environment {
environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
path: PathUri::parse("file:///ready-first-root")?,
},
};
let stale_root = SelectedCapabilityRoot {
id: refreshed_root.id.clone(),
location: CapabilityRootLocation::Environment {
environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
path: PathUri::parse("file:///stale-ready-first-root")?,
},
};
let expected_root = if executor_reports_refreshed_root {
refreshed_root.clone()
} else {
stale_root.clone()
};
let environment = test
.thread_manager
.environment_manager()
.report_environment_provisioning_status(
REMOTE_ENVIRONMENT_ID.to_string(),
Ok(EnvironmentReadyInfo {
selected_capability_roots: Vec::new(),
selected_capability_roots: if executor_reports_refreshed_root {
vec![refreshed_root.clone()]
} else {
Vec::new()
},
}),
Arc::new(ReadyNoiseConnectProvider {
websocket_url: format!("{rendezvous_url}/relay?role=harness"),
@@ -1867,26 +1899,45 @@ async fn ready_before_selection_exposes_remote_tools_and_capability_context_afte
anyhow::Ok(())
});
test.submit_turn_with_environments(
"use the ready environment",
Some(vec![TurnEnvironmentSelection {
environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
cwd: PathUri::from_abs_path(&test.config.cwd),
workspace_roots: vec![PathUri::from_abs_path(&test.config.cwd)],
config: EnvironmentConfigState::Ready(EnvironmentConfig {
allow_login_shell: true,
permission_profile: PermissionProfileSnapshot::legacy(
test.config.permissions.permission_profile().clone(),
),
shell_environment_policy: Default::default(),
exec_policy: None,
mcp_policy: None,
network_policy: None,
selected_capability_roots: vec![ready_root],
}),
}]),
)
.await?;
let selection = TurnEnvironmentSelection {
environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
cwd: PathUri::from_abs_path(&test.config.cwd),
workspace_roots: vec![PathUri::from_abs_path(&test.config.cwd)],
config: EnvironmentConfigState::FromThread,
};
let mut thread_extension_init = ExtensionDataInit::new();
thread_extension_init.insert(vec![stale_root]);
let resumed = test
.thread_manager
.start_thread(StartThreadOptions {
environments: Some(vec![selection.clone()]),
thread_extension_init,
..StartThreadOptions::new(test.config.clone())
})
.await?;
resumed
.thread
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
text: "use the ready environment".to_string(),
text_elements: Vec::new(),
}]))
.await?;
wait_for_event(&resumed.thread, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
assert_eq!(
resumed.thread.environment_selections().await,
vec![selection]
);
assert_eq!(
resumed
.thread
.inspect_selected_capability_roots()
.ready_roots,
vec![expected_root.clone()]
);
let requests = response_mock.requests();
assert_eq!(requests.len(), 2);
@@ -1930,6 +1981,14 @@ async fn ready_before_selection_exposes_remote_tools_and_capability_context_afte
.iter()
.any(|text| text.contains("<ready_capability_roots>ready-first-root"))
);
let observed_selected_roots = observed_roots
.lock()
.expect("observed capability roots should not be poisoned")
.iter()
.rfind(|roots| roots.iter().any(|root| root.id == expected_root.id))
.cloned()
.context("selected capability root should reach world state")?;
assert_eq!(observed_selected_roots, vec![expected_root]);
relay.abort();
remote_environment.abort();