Support pending environment attachment configuration (#38684)

## Why

An environment connection can be available before its owner has supplied the
configuration for a particular thread attachment. Threads need to start without
blocking while ensuring turns do not use that attachment prematurely.

## What changed

- Accept `Pending` environment configuration and resolve each attachment only
  after both its shared executor connection and owner configuration are ready.
- Add a `Failed` configuration state and `environment_failed` callback so an
  owner can fail one thread's attachment without affecting other threads.
- Keep pending and failed attachments out of capability-root inspection and turn
  environments, and allow failed attachments to recover through a ready update.
- Apply owner configuration before waking a waiting turn so its permission
  profile, login-shell policy, capability roots, and tools are immediately
  consistent.

## Testing

Add an integration test covering non-blocking thread startup, independent ready
and failed callbacks, waiting-turn resumption, installed capability and tool
configuration, and recovery from failure.

GitOrigin-RevId: d587e2025d584c867d782d470b18bf5a1a27b76c
This commit is contained in:
sayan-oai
2026-08-15 01:39:04 +00:00
committed by copyberry
parent eb147c0db3
commit 2ca575026c
7 changed files with 378 additions and 85 deletions
+9
View File
@@ -717,6 +717,15 @@ impl CodexThread {
self.session.environment_ready(selection, config).await
}
/// Fails this thread's pending environment without affecting other attached threads.
pub async fn environment_failed(
&self,
selection: &TurnEnvironmentSelection,
error: String,
) -> CodexResult<()> {
self.session.environment_failed(selection, error).await
}
/// Passively inspects the selected capability roots whose environments are ready now.
pub fn inspect_selected_capability_roots(&self) -> SelectedCapabilityRootsStatus {
self.session.inspect_selected_capability_roots()
+111 -54
View File
@@ -14,8 +14,6 @@ use codex_exec_server::ExecutorFileSystem;
use codex_exec_server::SelectedCapabilityRootsStatus;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::protocol::EnvironmentConfig;
use codex_protocol::protocol::EnvironmentConfigState;
use codex_protocol::protocol::EnvironmentConnectionEvent;
@@ -27,6 +25,7 @@ use codex_utils_path_uri::PathUri;
use futures::FutureExt;
use futures::future::BoxFuture;
use futures::future::Shared;
use tokio::sync::mpsc;
use tokio_util::task::AbortOnDropHandle;
use crate::session::turn_context::ShellSnapshotTask;
@@ -73,6 +72,7 @@ pub(crate) fn default_thread_environment_selections(
type TurnEnvironmentResult = Result<ResolvedEnvironment, Arc<ExecServerError>>;
type TurnEnvironmentResolution = Shared<BoxFuture<'static, TurnEnvironmentResult>>;
type PendingConfigurationResult = Result<EnvironmentConfig, String>;
// Shared startup result used to build each turn's environment with its own config
// without restarting the connection or shell resolution.
@@ -81,6 +81,7 @@ struct ResolvedEnvironment {
environment: Arc<Environment>,
shell: Option<Shell>,
shell_snapshot: ShellSnapshotTask,
installed_config: Option<EnvironmentConfig>,
}
#[derive(Clone)]
@@ -91,6 +92,7 @@ struct SelectedTurnEnvironment {
// Selection clones share one listener; the final handle drop aborts it.
connection_events_task: Option<Arc<AbortOnDropHandle<()>>>,
resolution: TurnEnvironmentResolution,
pending_completion: Option<mpsc::Sender<PendingConfigurationResult>>,
}
#[derive(Clone)]
@@ -106,15 +108,15 @@ fn resolve_selection_config(
thread_config: &EnvironmentConfig,
) -> (TurnEnvironmentSelection, EnvironmentConfigOrigin) {
let (config, origin) = match selection.config {
EnvironmentConfigState::FromThread => {
(thread_config.clone(), EnvironmentConfigOrigin::Thread)
}
EnvironmentConfigState::Ready(config) => (config, EnvironmentConfigOrigin::Owner),
EnvironmentConfigState::Pending => {
unreachable!("pending environment configuration is not supported yet")
}
EnvironmentConfigState::FromThread => (
EnvironmentConfigState::Ready(thread_config.clone()),
EnvironmentConfigOrigin::Thread,
),
config @ (EnvironmentConfigState::Ready(_)
| EnvironmentConfigState::Pending
| EnvironmentConfigState::Failed(_)) => (config, EnvironmentConfigOrigin::Owner),
};
selection.config = EnvironmentConfigState::Ready(config);
selection.config = config;
(selection, origin)
}
@@ -176,6 +178,7 @@ impl ThreadEnvironments {
environment: environment.environment,
shell: environment.shell,
shell_snapshot: environment.shell_snapshot,
installed_config: None,
}))
.boxed()
.shared();
@@ -185,6 +188,7 @@ impl ThreadEnvironments {
environment: selected_environment,
connection_events_task: None,
resolution,
pending_completion: None,
};
// Child threads re-infer thread-owned config while preserving owner config.
inherited_environment.refresh_thread_config(&thread_environment_config);
@@ -220,13 +224,26 @@ impl ThreadEnvironments {
previous.environment_id == selected_environment.environment_id
&& previous.cwd == selected_environment.cwd
&& previous.workspace_roots == selected_environment.workspace_roots
}) && !matches!(environment.resolution.clone().now_or_never(), Some(Err(_)))
{
let mut environment = environment.clone();
environment.selection = selected_environment;
environment.config_origin = config_origin;
next.push(environment);
continue;
}) {
let failed =
matches!(
environment.selection.config,
EnvironmentConfigState::Failed(_)
) || matches!(environment.resolution.clone().now_or_never(), Some(Err(_)));
let restarting_as_pending =
matches!(selected_environment.config, EnvironmentConfigState::Pending)
&& !matches!(
environment.selection.config,
EnvironmentConfigState::Pending
);
if !failed && !restarting_as_pending {
let mut environment = environment.clone();
environment.selection = selected_environment;
environment.config_origin = config_origin;
next.push(environment);
continue;
}
}
let environment_id = &selected_environment.environment_id;
@@ -251,11 +268,19 @@ impl ThreadEnvironments {
)
})
});
let (pending_completion, configuration_ready) =
if matches!(selected_environment.config, EnvironmentConfigState::Pending) {
let (sender, receiver) = mpsc::channel(/*buffer*/ 1);
(Some(sender), Some(receiver))
} else {
(None, None)
};
let (resolution_task, resolution) = Self::resolve_environment(
selected_environment.clone(),
Arc::clone(&environment),
self.local_shell.clone(),
self.shell_snapshot.clone(),
configuration_ready,
)
.remote_handle();
drop(tokio::spawn(resolution_task));
@@ -266,6 +291,7 @@ impl ThreadEnvironments {
environment,
connection_events_task,
resolution,
pending_completion,
};
next.push(selected);
}
@@ -281,7 +307,22 @@ impl ThreadEnvironments {
.then(|| Arc::clone(task))
})
.collect::<Vec<_>>();
self.environments.store(Arc::new(next));
let next = Arc::new(next);
self.environments.store(Arc::clone(&next));
// Publish owner configuration before waking turns waiting on this attachment.
for environment in next.iter() {
let Some(completion) = &environment.pending_completion else {
continue;
};
let result = match &environment.selection.config {
EnvironmentConfigState::Ready(config) => Ok(config.clone()),
EnvironmentConfigState::Failed(error) => Err(error.clone()),
EnvironmentConfigState::FromThread | EnvironmentConfigState::Pending => continue,
};
let _ = completion.try_send(result);
}
// ArcSwap readers may retain removed selections, so abort at logical removal.
for task in removed_connection_tasks {
task.abort();
@@ -340,30 +381,6 @@ impl ThreadEnvironments {
self.environments.store(Arc::new(environments));
}
/// Installs owner-provided config and roots on their exact thread attachment.
/// Additional environment-owned settings should be applied to its config here.
pub(crate) fn environment_ready(
&self,
selection: &TurnEnvironmentSelection,
config: EnvironmentConfig,
) -> CodexResult<()> {
let mut environments = Vec::clone(&self.environments.load());
let Some(environment) = environments.iter_mut().find(|environment| {
environment.selection.environment_id == selection.environment_id
&& environment.selection.cwd == selection.cwd
}) else {
return Err(CodexErr::InvalidRequest(format!(
"environment `{}` is not selected on this thread with the requested cwd",
selection.environment_id
)));
};
environment.selection.config = EnvironmentConfigState::Ready(config);
environment.config_origin = EnvironmentConfigOrigin::Owner;
self.environments.store(Arc::new(environments));
Ok(())
}
/// Combines persisted thread roots with installed attachment roots, keeping
/// thread roots first and hiding attachments that are not ready yet.
pub(crate) fn inspect_selected_capability_roots(
@@ -374,7 +391,7 @@ impl ThreadEnvironments {
let mut selected_capability_roots = thread_selected_capability_roots.to_vec();
for environment in environments.iter() {
let EnvironmentConfigState::Ready(config) = &environment.selection.config else {
unreachable!("selected environments always carry resolved configuration")
continue;
};
selected_capability_roots.extend(config.selected_capability_roots.iter().cloned());
}
@@ -465,13 +482,36 @@ impl ThreadEnvironments {
environment: Arc<Environment>,
local_shell: Shell,
shell_snapshot: ShellSnapshot,
configuration_ready: Option<mpsc::Receiver<PendingConfigurationResult>>,
) -> BoxFuture<'static, TurnEnvironmentResult> {
async move {
let environment_id = &selection.environment_id;
if let Err(err) = environment.wait_until_ready().await {
tracing::warn!("turn environment `{environment_id}` failed to start: {err}");
return Err(Arc::new(err));
if let EnvironmentConfigState::Failed(error) = &selection.config {
return Err(Arc::new(ExecServerError::Protocol(error.clone())));
}
// Wait for the shared executor connection.
let connection_ready = async {
environment.wait_until_ready().await.map_err(|error| {
tracing::warn!("turn environment `{environment_id}` failed to start: {error}");
Arc::new(error)
})
};
// Wait for this thread attachment's owner configuration, if pending.
let configuration_ready = async move {
let Some(mut configuration_ready) = configuration_ready else {
return Ok(None);
};
match configuration_ready.recv().await {
Some(Ok(config)) => Ok(Some(config)),
Some(Err(error)) => Err(Arc::new(ExecServerError::Protocol(error))),
None => Err(Arc::new(ExecServerError::Protocol(
"environment configuration was canceled".to_string(),
))),
}
};
// Resolve the attachment only after both prerequisites are ready.
let ((), installed_config) = tokio::try_join!(connection_ready, configuration_ready)?;
let shell = if environment.is_remote() {
match environment.info().await {
Ok(info) => match Shell::from_environment_shell_info(info.shell) {
@@ -502,6 +542,7 @@ impl ThreadEnvironments {
environment,
shell,
shell_snapshot: task,
installed_config,
})
}
.boxed()
@@ -512,7 +553,17 @@ impl ThreadEnvironments {
let selected = self.environments.load_full();
let mut environments = Vec::with_capacity(selected.len());
for environment in selected.iter() {
let resolved = if self.non_blocking_snapshots {
if matches!(
environment.selection.config,
EnvironmentConfigState::Failed(_)
) {
continue;
}
let pending = matches!(
environment.selection.config,
EnvironmentConfigState::Pending
);
let resolved = if self.non_blocking_snapshots || pending {
environment.resolution.clone().now_or_never()
} else {
Some(environment.resolution.clone().await)
@@ -549,8 +600,12 @@ impl TurnEnvironmentState {
) -> Option<Self> {
match resolved {
Some(Ok(environment)) => {
let mut selection = starting.selection;
if matches!(selection.config, EnvironmentConfigState::Pending) {
selection.config = EnvironmentConfigState::Ready(environment.installed_config?);
}
let mut turn_environment = TurnEnvironment::new(
starting.selection,
selection,
starting.config_origin,
environment.environment,
environment.shell,
@@ -1481,9 +1536,12 @@ url = "ws://127.0.0.1:8765"
permission_profile: PermissionProfileSnapshot::legacy(PermissionProfile::read_only()),
selected_capability_roots: vec![root("parent-root")],
};
parent
.environment_ready(&selection, parent_owner_config.clone())
.expect("install environment config");
let mut owner_selection = selection.clone();
owner_selection.config = EnvironmentConfigState::Ready(parent_owner_config.clone());
parent.update_selections(
std::slice::from_ref(&owner_selection),
&test_environment_config(),
);
let child_thread_config = EnvironmentConfig {
permission_profile: PermissionProfileSnapshot::legacy(
@@ -1516,9 +1574,8 @@ url = "ws://127.0.0.1:8765"
selected_capability_roots: Vec::new(),
..parent_owner_config.clone()
};
parent
.environment_ready(&selection, cleared_parent_owner_config.clone())
.expect("clear parent roots");
owner_selection.config = EnvironmentConfigState::Ready(cleared_parent_owner_config.clone());
parent.update_selections(std::slice::from_ref(&owner_selection), &child_thread_config);
let cleared_snapshot = parent.snapshot().await;
assert_eq!(
cleared_snapshot
+41 -12
View File
@@ -19,12 +19,9 @@ pub(super) fn validate_environment_selections(
) -> CodexResult<()> {
for selection in selections {
match &selection.config {
EnvironmentConfigState::FromThread => {}
EnvironmentConfigState::Pending => {
return Err(CodexErr::InvalidRequest(
"pending environment configuration is not supported yet".to_string(),
));
}
EnvironmentConfigState::FromThread
| EnvironmentConfigState::Pending
| EnvironmentConfigState::Failed(_) => {}
EnvironmentConfigState::Ready(config) => {
validate_environment_config(selection, config)?;
}
@@ -74,14 +71,46 @@ impl Session {
config: EnvironmentConfig,
) -> CodexResult<()> {
validate_environment_config(selection, &config)?;
self.update_environment_configuration(selection, EnvironmentConfigState::Ready(config))
.await
}
// grab session lock so installation can't race w/ thread settings updates
let _state = self.state.lock().await;
self.services
.turn_environments
.environment_ready(selection, config)?;
// mark mcp runtime for refresh because available capabilities could've changed
pub(crate) async fn environment_failed(
&self,
selection: &TurnEnvironmentSelection,
error: String,
) -> CodexResult<()> {
self.update_environment_configuration(selection, EnvironmentConfigState::Failed(error))
.await
}
async fn update_environment_configuration(
&self,
selection: &TurnEnvironmentSelection,
config: EnvironmentConfigState,
) -> CodexResult<()> {
// Serialize owner callbacks with ordinary thread settings updates.
let state = self.state.lock().await;
let mut environments = self.services.turn_environments.selections();
let Some(environment) = environments.iter_mut().find(|environment| {
environment.environment_id == selection.environment_id
&& environment.cwd == selection.cwd
&& environment.workspace_roots == selection.workspace_roots
}) else {
return Err(CodexErr::InvalidRequest(format!(
"environment `{}` is not selected on this thread with the requested workspace",
selection.environment_id
)));
};
environment.config = config;
// Invalidate MCP before installed configuration can wake a waiting turn.
self.mark_mcp_runtime_dirty();
self.services.turn_environments.update_selections(
&environments,
&state.session_configuration.turn_environment_config(),
);
Ok(())
}
+3 -3
View File
@@ -1588,6 +1588,9 @@ impl Session {
previous_permission_profile != updated_permission_profile;
let mcp_inputs_changed =
self.mcp_inputs_differ(&state.session_configuration, &updated, &updates);
if mcp_inputs_changed {
self.mark_mcp_runtime_dirty();
}
let environment_config = updated.turn_environment_config();
if let Some(environments) = &updates.environments {
self.services
@@ -1601,9 +1604,6 @@ impl Session {
state.session_configuration = updated;
let new_config = notify_config_contributors
.then(|| self.build_effective_session_config(&state.session_configuration));
if mcp_inputs_changed {
self.mark_mcp_runtime_dirty();
}
(
previous_config,
new_config,
+3 -3
View File
@@ -693,6 +693,9 @@ impl Session {
Ok(next) => {
let mcp_inputs_changed =
self.mcp_inputs_differ(&state.session_configuration, &next, &updates);
if mcp_inputs_changed {
self.mark_mcp_runtime_dirty();
}
let previous_permission_profile =
state.session_configuration.permission_profile();
let next_permission_profile = next.permission_profile();
@@ -712,9 +715,6 @@ impl Session {
.turn_environments
.update_thread_config(&environment_config);
}
if mcp_inputs_changed {
self.mark_mcp_runtime_dirty();
}
state.session_configuration = next.clone();
let new_config = notify_config_contributors
.then(|| self.build_effective_session_config(&state.session_configuration));
+209 -13
View File
@@ -1174,22 +1174,22 @@ async fn shared_executor_keeps_ready_capability_roots_scoped_to_each_attachment(
selected_capability_roots: vec![root("duplicate"), root("duplicate")],
}),
] {
let invalid_selection = TurnEnvironmentSelection {
let should_succeed = matches!(config, EnvironmentConfigState::Pending);
let selection_override = TurnEnvironmentSelection {
config,
..selection.clone()
};
assert!(
test.codex
.preview_thread_settings_overrides(CodexThreadSettingsOverrides {
environments: Some(TurnEnvironmentSelections::new(
test.config.cwd.clone(),
vec![invalid_selection],
)),
..Default::default()
})
.await
.is_err()
);
let preview = test
.codex
.preview_thread_settings_overrides(CodexThreadSettingsOverrides {
environments: Some(TurnEnvironmentSelections::new(
test.config.cwd.clone(),
vec![selection_override],
)),
..Default::default()
})
.await;
assert_eq!(preview.is_ok(), should_succeed);
assert_eq!(
test.codex.environment_selections().await,
vec![selection.clone()]
@@ -1334,6 +1334,202 @@ async fn shared_executor_keeps_ready_capability_roots_scoped_to_each_attachment(
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn pending_attachment_installs_configuration_before_waiting_turn_resumes() -> Result<()> {
const WAIT_CALL_ID: &str = "wait-for-owner-configuration";
let server = start_mock_server().await;
let mut extensions = ExtensionRegistryBuilder::new();
extensions.thread_lifecycle_contributor(Arc::new(WaitForEnvironmentTestExtension));
extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension));
let mut builder = test_codex()
.with_extensions(Arc::new(extensions.build()))
.with_config(|config| {
config.use_experimental_unified_exec_tool = true;
assert!(config.features.enable(Feature::DeferredExecutor).is_ok());
assert!(config.features.enable(Feature::UnifiedExec).is_ok());
config
.permissions
.set_permission_profile(PermissionProfile::read_only())
.expect("thread permissions should be configurable");
});
let test = builder.build_with_auto_env(&server).await?;
let selection = test
.codex
.environment_selections()
.await
.into_iter()
.next()
.context("thread should select its executor environment")?;
let pending_selection = TurnEnvironmentSelection {
config: EnvironmentConfigState::Pending,
..selection.clone()
};
let root = |id: &str| SelectedCapabilityRoot {
id: id.to_string(),
location: CapabilityRootLocation::Environment {
environment_id: selection.environment_id.clone(),
path: selection.cwd.clone(),
},
};
let owner_config = |id: &str, allow_login_shell: bool| EnvironmentConfig {
allow_login_shell,
permission_profile: PermissionProfileSnapshot::legacy(PermissionProfile::read_only()),
selected_capability_roots: vec![root(id)],
};
let start_pending_thread = || {
test.thread_manager.start_thread(StartThreadOptions {
environments: Some(vec![pending_selection.clone()]),
..StartThreadOptions::new(test.config.clone())
})
};
let waiting = timeout(Duration::from_secs(5), start_pending_thread())
.await
.context("pending thread startup should not block")??;
let independent = start_pending_thread().await?;
let failed = start_pending_thread().await?;
submit_thread_settings(
&waiting.thread,
ThreadSettingsOverrides {
permission_profile: Some(PermissionProfile::workspace_write()),
..Default::default()
},
)
.await?;
let response_mock = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("pending-configuration-wait"),
ev_function_call(
WAIT_CALL_ID,
"wait_for_environment",
&json!({ "environment_id": selection.environment_id }).to_string(),
),
ev_completed("pending-configuration-wait"),
]),
sse(vec![
ev_response_created("pending-configuration-ready"),
ev_assistant_message("pending-configuration-message", "done"),
ev_completed("pending-configuration-ready"),
]),
],
)
.await;
waiting
.thread
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
text: "wait for environment configuration".into(),
text_elements: Vec::new(),
}]))
.await?;
wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await;
let first_tool_names = tool_names(&response_mock.requests()[0].body_json());
assert!(first_tool_names.contains(&"wait_for_environment".to_string()));
assert!(!first_tool_names.contains(&"exec_command".to_string()));
independent
.thread
.environment_ready(
&pending_selection,
owner_config("independent-root", /*allow_login_shell*/ true),
)
.await?;
failed
.thread
.environment_failed(&pending_selection, "configuration unavailable".to_string())
.await?;
assert_eq!(
failed.thread.environment_selections().await,
vec![TurnEnvironmentSelection {
config: EnvironmentConfigState::Failed("configuration unavailable".to_string()),
..pending_selection.clone()
}]
);
assert!(
waiting
.thread
.inspect_selected_capability_roots()
.ready_roots
.is_empty()
);
let waiting_config = owner_config("waiting-root", /*allow_login_shell*/ false);
let ready_selection = TurnEnvironmentSelection {
config: EnvironmentConfigState::Ready(waiting_config.clone()),
..pending_selection.clone()
};
submit_thread_settings(
&waiting.thread,
ThreadSettingsOverrides {
environments: Some(TurnEnvironmentSelections::new(
test.config.cwd.clone(),
vec![ready_selection.clone()],
)),
..Default::default()
},
)
.await?;
wait_for_response_request_count(&response_mock, /*expected_count*/ 2).await;
assert_eq!(
waiting.thread.environment_selections().await,
vec![ready_selection]
);
let ready_request = response_mock
.last_request()
.context("waiting turn should resume")?;
let (_, wait_succeeded) = ready_request
.function_call_output_content_and_success(WAIT_CALL_ID)
.context("wait_for_environment output should be model visible")?;
assert_ne!(wait_succeeded, Some(false));
let body = ready_request.body_json();
let exec_command = body["tools"]
.as_array()
.context("tools should be an array")?
.iter()
.find(|tool| tool["name"] == "exec_command")
.context("exec_command should become available")?;
assert!(
exec_command["parameters"]["properties"]
.get("login")
.is_none()
);
assert_eq!(
ready_request
.message_input_texts("user")
.into_iter()
.rfind(|text| text.contains("<ready_capability_roots>")),
Some("<ready_capability_roots>waiting-root</ready_capability_roots>".to_string())
);
let recovered_selection = TurnEnvironmentSelection {
config: EnvironmentConfigState::Ready(owner_config(
"recovered-root",
/*allow_login_shell*/ true,
)),
..pending_selection
};
submit_thread_settings(
&failed.thread,
ThreadSettingsOverrides {
environments: Some(TurnEnvironmentSelections::new(
test.config.cwd.clone(),
vec![recovered_selection.clone()],
)),
..Default::default()
},
)
.await?;
assert_eq!(
failed.thread.environment_selections().await,
vec![recovered_selection]
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn ready_before_selection_exposes_remote_tools_and_capability_context_after_wait()
-> Result<()> {
+2
View File
@@ -10,6 +10,8 @@ pub enum EnvironmentConfigState {
Pending,
/// The owner supplied configuration for this environment attachment.
Ready(EnvironmentConfig),
/// The owner could not supply configuration for this environment attachment.
Failed(String),
}
/// Resolved configuration for a thread/environment attachment.