Defer environment selection changes until the next turn (#46310)

## Why

Updating environment selections while a turn is running must not redirect its tools or prevent its pending environment setup from completing.

## What changed

- Store future environment selections separately from active environments, and activate them when preparing new work with no task running.
- Keep turn contexts and MCP workspace roots tied to the active environment snapshot. Report saved selections through thread settings and `environment_selections()`.
- Route environment configuration and failure callbacks to matching active or future selections so each can finish setup independently.
- Stop the running task before manual compaction adopts the saved environments.

## Testing

Add regression tests that preserve a running task's working directory and workspace roots, then adopt the saved selection for the next task. Add remote environment coverage showing that future setup can complete while the active turn waits, with each turn receiving its own capability roots.

GitOrigin-RevId: 09a6b90131d2ab2441ce2fe9b0368a3668889f31
This commit is contained in:
sayan-oai
2026-09-17 21:47:23 +00:00
committed by copyberry
parent 0c9be8a836
commit c775dd3c33
15 changed files with 522 additions and 135 deletions
+18 -13
View File
@@ -322,18 +322,13 @@ impl AgentControl {
parent: Option<Arc<CodexThread>>,
) -> CodexResult<()> {
let state = self.upgrade()?;
let parent = if let Some(parent) = parent {
let owner_thread_id = parent.as_ref().map(|parent| parent.session.thread_id);
if let Some(parent) = &parent {
let parent_thread_id = parent.session.thread_id;
let turn = parent.session.new_default_turn().await;
config = build_agent_resume_config(&turn).map_err(|_| {
CodexErr::InvalidRequest(format!(
"cannot resume multi-agent v2 child {thread_id} with the current parent settings"
))
})?;
let registered_parent = state.get_thread(parent_thread_id).await.ok();
if !registered_parent
.as_ref()
.is_some_and(|registered| Arc::ptr_eq(registered, &parent))
.is_some_and(|registered| Arc::ptr_eq(registered, parent))
|| !parent.is_running()
|| parent.multi_agent_version() != Some(MultiAgentVersion::V2)
|| !Arc::ptr_eq(&self.state, &parent.session.services.agent_control.state)
@@ -342,11 +337,7 @@ impl AgentControl {
"cannot resume multi-agent v2 child {thread_id}: parent ownership is unavailable; resume the parent first"
)));
}
Some((parent, turn.environments.clone()))
} else {
None
};
let owner_thread_id = parent.as_ref().map(|(parent, _)| parent.session.thread_id);
}
if owner_thread_id.is_none() && state.get_thread(thread_id).await.is_ok() {
self.touch_loaded_v2_residency(&state, thread_id).await;
return Ok(());
@@ -400,6 +391,20 @@ impl AgentControl {
return Ok(());
}
}
let parent = if let Some(parent) = parent {
let turn = parent
.session
.new_turn_with_default_settings(Uuid::now_v7().to_string(), Default::default())
.await;
config = build_agent_resume_config(&turn).map_err(|_| {
CodexErr::InvalidRequest(format!(
"cannot resume multi-agent v2 child {thread_id} with the current parent settings"
))
})?;
Some((parent, turn.environments.clone()))
} else {
None
};
config.model_reasoning_effort = stored_reasoning_effort;
if let Some(role_name) = session_source.get_agent_role() {
let runtime_approval_policy = config.permissions.approval_policy.value();
+13 -6
View File
@@ -14,6 +14,7 @@ use crate::context::ManagedDeveloperInstructions;
use crate::context::MultiAgentRoleInstructions;
use crate::context::SubagentNotification;
use crate::init_state_db;
use crate::session::SessionSettingsUpdate;
use crate::thread_manager::StartThreadOptions;
use crate::tools::handlers::multi_agents_common::thread_spawn_source;
use assert_matches::assert_matches;
@@ -66,6 +67,7 @@ use codex_protocol::protocol::TokenUsageRecord;
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnAbortedEvent;
use codex_protocol::protocol::TurnCompleteEvent;
use codex_protocol::protocol::TurnEnvironmentSelections;
use codex_protocol::protocol::TurnStartedEvent;
use codex_thread_store::ArchiveThreadParams;
use codex_thread_store::InMemoryThreadStore;
@@ -897,17 +899,21 @@ async fn check_v2_agent_reload(route: V2ReloadRoute) {
.environments
.primary()
.expect("parent environment");
let thread_config = environment.config().clone();
let mut owner_config = thread_config.clone();
let mut owner_config = environment.config().clone();
owner_config.allow_login_shell = false;
let mut selection = environment.selection();
selection.config = EnvironmentConfigState::Ready(owner_config);
parent_thread
.session
.services
.turn_environments
.update_selections(std::slice::from_ref(&selection), &thread_config);
parent_turn = parent_thread.session.new_default_turn().await;
.update_settings(SessionSettingsUpdate {
environments: Some(TurnEnvironmentSelections::new(
parent_turn.config.cwd.clone(),
vec![selection],
)),
..Default::default()
})
.await
.expect("save parent environments");
parent_thread.session.mark_interrupted();
// The fixture has no task runner to finish the turn or consume child results.
*parent_thread.session.active_turn.lock().await = None;
@@ -921,6 +927,7 @@ async fn check_v2_agent_reload(route: V2ReloadRoute) {
.ensure_multi_agent_v2_child_loaded(spawned_agent.thread_id)
.await
.expect("known child should reload through its parent");
parent_turn = parent_thread.session.new_default_turn().await;
assert!(harness.manager.get_thread(parent_thread_id).await.is_err());
}
}
+2 -1
View File
@@ -908,8 +908,9 @@ impl CodexThread {
self.session.refresh_codex_apps_tools().await
}
/// Returns the environments configured for future turns.
pub async fn environment_selections(&self) -> Vec<TurnEnvironmentSelection> {
self.session.services.turn_environments.selections()
self.session.configured_environment_selections().await
}
/// Installs resolved environment configuration and capability roots on this thread.
+49 -38
View File
@@ -1,6 +1,7 @@
use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt;
use std::future::Future;
use std::sync::Arc;
use std::sync::OnceLock;
@@ -490,15 +491,6 @@ impl ThreadEnvironments {
.collect()
}
pub(crate) fn primary_workspace_roots(&self) -> Vec<AbsolutePathBuf> {
self.environments
.load()
.first()
.map_or_else(Vec::new, |environment| {
Self::primary_workspace_roots_for(std::slice::from_ref(&environment.selection))
})
}
/// Returns installed owner configuration without treating pending attachments as ready.
pub(crate) fn primary_config_for(
selections: &[TurnEnvironmentSelection],
@@ -806,45 +798,53 @@ impl ThreadEnvironments {
})
}
/// Captures the selected list immediately, then returns a future that may wait for setup.
/// This lets callers release their locks before waiting without accidentally using a newer
/// selection. The trace still covers the setup wait itself.
#[tracing::instrument(
name = "environments.snapshot",
skip_all,
fields(
environment_count = self.environments.load().len(),
non_blocking = self.non_blocking_snapshots,
environment_count = selected.len(),
non_blocking = non_blocking_snapshots,
)
)]
pub(crate) async fn snapshot(&self) -> TurnEnvironmentSnapshot {
pub(crate) fn snapshot(
&self,
) -> impl Future<Output = TurnEnvironmentSnapshot> + Send + 'static {
let selected = self.environments.load_full();
let mut environments = Vec::with_capacity(selected.len());
for environment in selected.iter() {
if let EnvironmentConfigState::Failed(error) = &environment.selection.config {
environments.push(TurnEnvironmentState::Failed {
let non_blocking_snapshots = self.non_blocking_snapshots;
async move {
let mut environments = Vec::with_capacity(selected.len());
for environment in selected.iter() {
if let EnvironmentConfigState::Failed(error) = &environment.selection.config {
environments.push(TurnEnvironmentState::Failed {
selection: environment.selection.clone(),
error: error.clone(),
});
continue;
}
let pending = matches!(
environment.selection.config,
EnvironmentConfigState::Pending
);
let starting = StartingTurnEnvironment {
selection: environment.selection.clone(),
error: error.clone(),
});
continue;
config_origin: environment.config_origin,
resolution: environment.resolution.clone(),
};
let resolved = if non_blocking_snapshots || pending {
starting.resolution.clone().now_or_never()
} else {
Some(match starting.wait_until_ready().await {
Ok(()) => starting.resolution.clone().await,
Err(error) => Err(error),
})
};
environments.push(TurnEnvironmentState::from_resolution(starting, resolved));
}
let pending = matches!(
environment.selection.config,
EnvironmentConfigState::Pending
);
let starting = StartingTurnEnvironment {
selection: environment.selection.clone(),
config_origin: environment.config_origin,
resolution: environment.resolution.clone(),
};
let resolved = if self.non_blocking_snapshots || pending {
starting.resolution.clone().now_or_never()
} else {
Some(match starting.wait_until_ready().await {
Ok(()) => starting.resolution.clone().await,
Err(error) => Err(error),
})
};
environments.push(TurnEnvironmentState::from_resolution(starting, resolved));
TurnEnvironmentSnapshot { environments }
}
TurnEnvironmentSnapshot { environments }
}
pub(crate) fn environment_manager(&self) -> Arc<EnvironmentManager> {
@@ -1000,6 +1000,17 @@ impl TurnEnvironmentSnapshot {
self.turn_environments().next()
}
/// Returns the first selected environment's host folders, even if setup is not ready yet.
pub(crate) fn primary_workspace_roots(&self) -> Vec<AbsolutePathBuf> {
let selection = match self.environments.first() {
Some(TurnEnvironmentState::Ready(environment)) => &environment.selection,
Some(TurnEnvironmentState::Starting(environment)) => &environment.selection,
Some(TurnEnvironmentState::Failed { selection, .. }) => selection,
None => return Vec::new(),
};
ThreadEnvironments::primary_workspace_roots_for(std::slice::from_ref(selection))
}
/// Returns the primary environment's resolved permissions, or the provided fallback.
pub(crate) fn permission_profile_or_else(
&self,
+103 -22
View File
@@ -14,6 +14,7 @@ use codex_protocol::protocol::TurnEnvironmentSelection;
use crate::config::ConstraintError;
use crate::config::ConstraintResult;
use crate::config::NetworkProxySpec;
use crate::environment_selection::TurnEnvironmentSnapshot;
use crate::session::session::Session;
use crate::session::session::SessionConfiguration;
use crate::session::session::SessionSettingsUpdate;
@@ -108,7 +109,7 @@ impl Session {
current: &SessionConfiguration,
updates: &SessionSettingsUpdate,
) -> ConstraintResult<SessionConfiguration> {
let current_environments = self.services.turn_environments.selections();
let current_environments = &current.environments;
if let Some(environments) = &updates.environments
&& let Some(environment) = environments.environments.iter().find(|environment| {
environment.config == EnvironmentConfigState::FromThread
@@ -126,7 +127,47 @@ impl Session {
});
}
current.apply(updates, &current_environments)
current.apply(updates, current_environments)
}
/// Activates the environments accepted for a new task. Configuration may have arrived for
/// those same environments while its settings were being prepared.
#[expect(
clippy::await_holding_invalid_type,
reason = "do not change environments while another task is running"
)]
pub(super) async fn activate_turn_environments(
&self,
configuration: &SessionConfiguration,
) -> TurnEnvironmentSnapshot {
let snapshot = {
let active = self.active_turn.lock().await;
let state = self.state.lock().await;
// Another task may have started while this one was preparing its settings.
if active.as_ref().is_none_or(|turn| turn.task.is_none()) {
let mut environments = configuration.environments.clone();
let latest_environments = &state.session_configuration.environments;
for environment in &mut environments {
if let Some(latest) = latest_environments.iter().find(|latest| {
latest.environment_id == environment.environment_id
&& latest.cwd == environment.cwd
&& latest.workspace_roots == environment.workspace_roots
}) {
environment.config = latest.config.clone();
}
}
if environments != self.services.turn_environments.selections() {
self.mark_mcp_runtime_dirty();
self.services.turn_environments.update_selections(
&environments,
&configuration.inferred_environment_config(),
);
self.schedule_mcp_prewarm();
}
}
self.services.turn_environments.snapshot()
};
snapshot.await
}
pub(crate) async fn environment_ready(
@@ -148,39 +189,79 @@ impl Session {
.await
}
/// Records configuration, or a failure to obtain it, for this environment and workspace.
/// The running turn and future turns may use different environments. Update the matching
/// selection or selections so work already waiting can still finish, and check successful
/// configuration before saving either list.
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| {
let mut state = self.state.lock().await;
let mut current = self.services.turn_environments.selections();
let mut future = state.session_configuration.environments.clone();
let matches = |environment: &TurnEnvironmentSelection| {
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
)));
};
let (current_environment, future_environment) = match (
current.iter_mut().find(|environment| matches(environment)),
future.iter_mut().find(|environment| matches(environment)),
) {
(None, None) => {
return Err(CodexErr::InvalidRequest(format!(
"environment `{}` is not selected on this thread with the requested workspace",
selection.environment_id
)));
}
(Some(current), Some(future)) if current.config != future.config => {
// Update the selection still waiting for configuration. If neither is waiting,
// keep the existing behavior of updating the current one.
if matches!(future.config, EnvironmentConfigState::Pending) {
(None, Some(future))
} else {
(Some(current), None)
}
}
(Some(current), future) => (Some(current), future),
(None, Some(future)) => (None, Some(future)),
};
let update_current = current_environment.is_some();
let update_future = future_environment.is_some();
environment.config = config;
if matches!(environment.config, EnvironmentConfigState::Ready(_)) {
state
.session_configuration
.validate(&environments)
.map_err(|error| CodexErr::InvalidRequest(error.to_string()))?;
if let Some(environment) = current_environment {
environment.config = config.clone();
}
if let Some(environment) = future_environment {
environment.config = config.clone();
}
if matches!(config, EnvironmentConfigState::Ready(_)) {
let validate = |environments: &[TurnEnvironmentSelection]| {
state
.session_configuration
.validate(environments)
.map_err(|error| CodexErr::InvalidRequest(error.to_string()))
};
if update_current {
validate(&current)?;
}
if update_future && (!update_current || current != future) {
validate(&future)?;
}
}
// 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.inferred_environment_config(),
);
state.session_configuration.environments = future;
if update_current {
// Invalidate MCP before installed configuration can wake a waiting turn.
self.mark_mcp_runtime_dirty();
self.services.turn_environments.update_selections(
&current,
&state.session_configuration.inferred_environment_config(),
);
}
Ok(())
}
+2
View File
@@ -241,6 +241,8 @@ pub async fn reload_user_config(sess: &Arc<Session>) {
}
pub async fn compact(sess: &Arc<Session>, sub_id: String) {
// Stop the old turn before the compact task picks up the next turn's environments.
sess.abort_all_tasks(TurnAbortReason::Replaced).await;
let turn_context = sess
.new_turn_with_default_settings(sub_id, Default::default())
.await;
+14 -8
View File
@@ -36,7 +36,6 @@ impl Session {
&self,
current: &SessionConfiguration,
next: &SessionConfiguration,
updates: &SessionSettingsUpdate,
) -> bool {
current.cwd() != next.cwd()
|| current.step_settings.approval_policy.value()
@@ -44,9 +43,6 @@ impl Session {
|| current.step_settings.approvals_reviewer != next.step_settings.approvals_reviewer
|| current.permission_profile() != next.permission_profile()
|| current.windows_sandbox_level != next.windows_sandbox_level
|| updates.environments.as_ref().is_some_and(|environments| {
environments.environments != self.services.turn_environments.selections()
})
}
/// Waits on this session's refreshed server before tool execution is admitted.
@@ -76,19 +72,25 @@ impl Session {
&self,
auth: Option<CodexAuth>,
) -> McpDesiredState {
let (session_configuration, disabled_plugin_ids) = {
let (session_configuration, disabled_plugin_ids, environments) = {
let state = self.state.lock().await;
// MCP tools must use current environments, not the selection saved for the next turn.
(
state.session_configuration.clone(),
state.active_disabled_plugin_ids.clone(),
self.services.turn_environments.snapshot(),
)
};
let environments = self.services.turn_environments.snapshot().await;
let environments = environments.await;
let cwd = environments
.primary()
.and_then(|environment| environment.cwd().to_abs_path().ok())
.unwrap_or_else(|| session_configuration.cwd().clone());
let config = self.build_per_turn_config(&session_configuration, cwd);
let config = self.build_per_turn_config(
&session_configuration,
cwd,
environments.primary_workspace_roots(),
);
let local_process_cwd = environments
.local_environment_cwd()
.unwrap_or_else(|| session_configuration.cwd().clone())
@@ -116,7 +118,11 @@ impl Session {
) -> anyhow::Result<()> {
let cwd = AbsolutePathBuf::from_absolute_path(mcp_runtime_cwd)
.unwrap_or_else(|_| session_configuration.cwd().clone());
let config = self.build_per_turn_config(session_configuration, cwd);
let config = self.build_per_turn_config(
session_configuration,
cwd,
resolved_environments.primary_workspace_roots(),
);
let local_process_cwd = resolved_environments
.local_environment_cwd()
.unwrap_or_else(|| session_configuration.cwd().clone())
+14 -19
View File
@@ -835,6 +835,7 @@ impl Session {
approval_policy: config.permissions.approval_policy.clone(),
approvals_reviewer: config.approvals_reviewer,
}),
environments: environment_selections.clone(),
model_info_overrides: config.to_models_manager_config().into(),
developer_instructions: config.developer_instructions.clone(),
base_instructions,
@@ -1882,22 +1883,15 @@ impl Session {
let updated_permission_profile = updated.permission_profile();
let permission_profile_changed =
previous_permission_profile != updated_permission_profile;
let mcp_inputs_changed =
self.mcp_inputs_differ(&state.session_configuration, &updated, &updates);
let root_service_tier_changed = updated.parent_thread_id.is_none()
&& state.session_configuration.step_settings.service_tier
!= updated.step_settings.service_tier;
let environment_config = updated.inferred_environment_config();
let mcp_inputs_changed = self.mcp_inputs_differ(&state.session_configuration, &updated);
if mcp_inputs_changed {
self.mark_mcp_runtime_dirty();
}
let environment_config = updated.inferred_environment_config();
if let Some(environments) = &updates.environments {
self.services
.turn_environments
.update_selections(&environments.environments, &environment_config);
} else if state.session_configuration.inferred_environment_config()
!= environment_config
{
if state.session_configuration.inferred_environment_config() != environment_config {
self.services
.turn_environments
.update_thread_config(&environment_config);
@@ -1918,7 +1912,7 @@ impl Session {
configuration: state.session_configuration.clone(),
snapshot: state
.session_configuration
.thread_settings_snapshot(&self.services.turn_environments.selections()),
.thread_settings_snapshot(&state.session_configuration.environments),
};
(
commit,
@@ -1945,32 +1939,33 @@ impl Session {
) -> ConstraintResult<ThreadConfigSnapshot> {
let state = self.state.lock().await;
let configuration = self.apply_session_settings(&state.session_configuration, updates)?;
let environments = updates.environments.as_ref().map_or_else(
|| self.services.turn_environments.selections(),
|environments| environments.environments.clone(),
);
Ok(configuration.thread_config_snapshot(environments))
Ok(configuration.thread_config_snapshot(configuration.environments.clone()))
}
pub(crate) async fn thread_config_snapshot(&self) -> ThreadConfigSnapshot {
let state = self.state.lock().await;
state
.session_configuration
.thread_config_snapshot(self.services.turn_environments.selections())
.thread_config_snapshot(state.session_configuration.environments.clone())
}
pub(crate) async fn configured_environment_selections(&self) -> Vec<TurnEnvironmentSelection> {
let state = self.state.lock().await;
state.session_configuration.environments.clone()
}
pub(crate) async fn thread_settings_snapshot(&self) -> ThreadSettingsSnapshot {
let state = self.state.lock().await;
state
.session_configuration
.thread_settings_snapshot(&self.services.turn_environments.selections())
.thread_settings_snapshot(&state.session_configuration.environments)
}
pub(crate) async fn restorable_thread_settings(&self) -> CodexThreadSettingsOverrides {
let state = self.state.lock().await;
state
.session_configuration
.restorable_thread_settings(self.services.turn_environments.selections())
.restorable_thread_settings(state.session_configuration.environments.clone())
}
pub(crate) async fn set_app_server_client_info(
+4
View File
@@ -104,6 +104,8 @@ pub(crate) struct SessionConfiguration {
/// Desired configured inputs inherited by future turns.
pub(super) step_settings: Arc<StepSettings>,
/// Ordered environments inherited by future turns.
pub(super) environments: Vec<TurnEnvironmentSelection>,
/// Explicit startup overrides used when resolving effective model metadata.
pub(super) model_info_overrides: ModelInfoOverrides,
@@ -419,6 +421,7 @@ impl SessionConfiguration {
let current_cwd = self.cwd().clone();
if let Some(environments) = &updates.environments {
next_configuration.legacy_fallback_cwd = environments.legacy_fallback_cwd.clone();
next_configuration.environments = environments.environments.clone();
}
let cwd_changed = next_configuration.legacy_fallback_cwd != current_cwd;
if let Some(runtime_workspace_roots) = &updates.runtime_workspace_roots {
@@ -1402,6 +1405,7 @@ impl Session {
environment_selections,
&session_configuration.inferred_environment_config(),
);
session_configuration.environments = turn_environments.selections();
let resolved_environments = turn_environments.snapshot().await;
let agents_md_manager = Arc::new(AgentsMdManager::new(instructions));
let plugin_skill_warmup = warm_plugins_and_skills_for_session_init(
+80 -10
View File
@@ -4014,6 +4014,7 @@ async fn set_rate_limits_retains_previous_credits() {
};
let session_configuration = SessionConfiguration {
provider: create_model_provider(config.model_provider.clone(), /*auth_manager*/ None),
environments: Vec::new(),
step_settings: Arc::new(StepSettings {
collaboration_mode,
reasoning_summary: config.model_reasoning_summary,
@@ -4136,6 +4137,7 @@ async fn set_rate_limits_updates_plan_type_when_present() {
};
let session_configuration = SessionConfiguration {
provider: create_model_provider(config.model_provider.clone(), /*auth_manager*/ None),
environments: Vec::new(),
step_settings: Arc::new(StepSettings {
collaboration_mode,
reasoning_summary: config.model_reasoning_summary,
@@ -4755,6 +4757,7 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati
SessionConfiguration {
provider: create_model_provider(config.model_provider.clone(), /*auth_manager*/ None),
environments: Vec::new(),
step_settings: Arc::new(StepSettings {
collaboration_mode,
reasoning_summary: config.model_reasoning_summary,
@@ -5742,7 +5745,7 @@ async fn relative_cwd_update_without_environments_resolves_under_session_cwd() {
let state = session.state.lock().await;
assert_eq!(state.session_configuration.cwd(), &updated_cwd);
assert!(session.services.turn_environments.selections().is_empty());
assert!(state.session_configuration.environments.is_empty());
}
#[tokio::test]
@@ -5772,7 +5775,7 @@ async fn environment_settings_preserve_explicit_primary_cwd() {
let state = session.state.lock().await;
assert_eq!(state.session_configuration.cwd(), &updated_cwd);
assert_eq!(
session.services.turn_environments.selections()[0].cwd,
state.session_configuration.environments[0].cwd,
PathUri::from_abs_path(&environment_cwd)
);
}
@@ -5841,6 +5844,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() {
config.model_provider.clone(),
Some(Arc::clone(&auth_manager)),
),
environments: Vec::new(),
step_settings: Arc::new(StepSettings {
collaboration_mode,
reasoning_summary: config.model_reasoning_summary,
@@ -6018,6 +6022,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
};
let default_environments = vec![local(config.cwd.clone())];
let session_configuration = SessionConfiguration {
environments: default_environments.clone(),
provider: create_model_provider(
config.model_provider.clone(),
Some(Arc::clone(&auth_manager)),
@@ -6234,8 +6239,13 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
forked_from_ordinal_exclusive: None,
next_internal_sub_id: AtomicU64::new(0),
};
let per_turn_config =
session.build_per_turn_config(&session_configuration, session_configuration.cwd().clone());
let per_turn_config = session.build_per_turn_config(
&session_configuration,
session_configuration.cwd().clone(),
crate::environment_selection::ThreadEnvironments::primary_workspace_roots_for(
&session_configuration.environments,
),
);
let plugins_input = per_turn_config.plugins_config_input();
let plugin_outcome = session
.services
@@ -6328,6 +6338,7 @@ async fn make_session_with_config_and_rx(
};
let default_environments = vec![local(config.cwd.clone())];
let session_configuration = SessionConfiguration {
environments: default_environments.clone(),
provider: create_model_provider(
config.model_provider.clone(),
Some(Arc::clone(&auth_manager)),
@@ -6460,6 +6471,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx(
};
let default_environments = vec![local(config.cwd.clone())];
let session_configuration = SessionConfiguration {
environments: default_environments.clone(),
provider: create_model_provider(
config.model_provider.clone(),
Some(Arc::clone(&auth_manager)),
@@ -7536,7 +7548,55 @@ async fn turn_environments_set_primary_environment() {
}
#[tokio::test]
async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_environments() {
async fn new_task_does_not_change_an_already_running_tasks_environments() {
let (session, running_turn, events) = make_session_and_context_with_rx().await;
let current_cwd = running_turn.config.cwd.clone();
let next_workspace = tempfile::tempdir().expect("next workspace");
let next_cwd = next_workspace.path().abs();
let finish = Arc::new(Notify::new());
session
.update_settings(SessionSettingsUpdate {
environments: Some(TurnEnvironmentSelections::new(
next_cwd.clone(),
vec![local(next_cwd.clone())],
)),
..Default::default()
})
.await
.expect("save next task's settings");
session
.spawn_task(
running_turn,
Vec::new(),
HeldStepTask {
kind: TaskKind::Regular,
finish: Arc::clone(&finish),
},
)
.await;
let candidate = session
.new_turn_with_default_settings("candidate".into(), Default::default())
.await;
assert_eq!(candidate.config.cwd, current_cwd);
assert_eq!(candidate.config.workspace_roots, vec![current_cwd.clone()]);
assert_eq!(
session.services.turn_environments.selections(),
vec![local(current_cwd)]
);
finish.notify_one();
recv_terminal_event(&events, TerminalEventKind::TurnComplete).await;
let next = session
.new_turn_with_default_settings("next".into(), Default::default())
.await;
assert_eq!(next.config.cwd, next_cwd);
assert_eq!(next.config.workspace_roots, vec![next_cwd]);
}
#[tokio::test]
async fn task_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_environments() {
let (session, _initial_turn, _rx) = make_session_and_context_with_rx().await;
let session_cwd = session.get_config().await.cwd.clone();
let selected_cwd =
@@ -7551,7 +7611,9 @@ async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_en
})
.await
.expect("environment selection update should succeed");
let turn_context = session.new_default_turn().await;
let turn_context = session
.new_turn_with_default_settings("task".into(), Default::default())
.await;
let turn_environments = &turn_context.environments;
assert_eq!(turn_environments.turn_environments().count(), 1);
@@ -7573,7 +7635,7 @@ async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_en
}
#[tokio::test]
async fn default_turn_honors_empty_stored_thread_environments() {
async fn task_turn_honors_empty_stored_thread_environments() {
let (session, _initial_turn, _rx) = make_session_and_context_with_rx().await;
let session_cwd = session.get_config().await.cwd.clone();
@@ -7587,7 +7649,9 @@ async fn default_turn_honors_empty_stored_thread_environments() {
})
.await
.expect("environment selection update should succeed");
let turn_context = session.new_default_turn().await;
let turn_context = session
.new_turn_with_default_settings("task".into(), Default::default())
.await;
assert!(turn_context.environments.primary().is_none());
assert!(
@@ -8202,6 +8266,7 @@ where
};
let default_environments = vec![local(config.cwd.clone())];
let session_configuration = SessionConfiguration {
environments: default_environments.clone(),
provider: create_model_provider(
config.model_provider.clone(),
Some(Arc::clone(&auth_manager)),
@@ -8421,8 +8486,13 @@ where
forked_from_ordinal_exclusive: None,
next_internal_sub_id: AtomicU64::new(0),
});
let per_turn_config =
session.build_per_turn_config(&session_configuration, session_configuration.cwd().clone());
let per_turn_config = session.build_per_turn_config(
&session_configuration,
session_configuration.cwd().clone(),
crate::environment_selection::ThreadEnvironments::primary_workspace_roots_for(
&session_configuration.environments,
),
);
let plugins_input = per_turn_config.plugins_config_input();
let plugin_outcome = session
.services
+34 -9
View File
@@ -769,6 +769,7 @@ impl Session {
&self,
session_configuration: &SessionConfiguration,
cwd: AbsolutePathBuf,
workspace_roots: Vec<AbsolutePathBuf>,
) -> Config {
// todo(aibrahim): store this state somewhere else so we don't need to mut config
let config = session_configuration.original_config_do_not_use.clone();
@@ -776,7 +777,6 @@ impl Session {
per_turn_config.cwd = cwd;
per_turn_config.permissions.approval_policy =
session_configuration.step_settings.approval_policy.clone();
let workspace_roots = self.services.turn_environments.primary_workspace_roots();
per_turn_config.workspace_roots = workspace_roots.clone();
per_turn_config
.permissions
@@ -818,8 +818,15 @@ impl Session {
&self,
session_configuration: &SessionConfiguration,
) -> Config {
let mut config =
self.build_per_turn_config(session_configuration, session_configuration.cwd().clone());
let workspace_roots =
crate::environment_selection::ThreadEnvironments::primary_workspace_roots_for(
&session_configuration.environments,
);
let mut config = self.build_per_turn_config(
session_configuration,
session_configuration.cwd().clone(),
workspace_roots,
);
config.model = Some(
session_configuration
.step_settings
@@ -1008,22 +1015,25 @@ impl Session {
if let Some(service_tier) = service_tier_for_turn {
Arc::make_mut(&mut configuration.step_settings).service_tier = Some(service_tier);
}
let turn_environments = self.activate_turn_environments(&configuration).await;
let turn_context = self
.new_turn_from_configuration(sub_id, configuration, options)
.new_turn_from_configuration(sub_id, configuration, turn_environments, options)
.await;
Ok(Some((turn_context, commit.snapshot)))
}
/// Constructs a turn from the exact committed settings without starting a task.
/// Builds a context from the caller's chosen settings and environments without starting work.
async fn new_turn_from_configuration(
&self,
sub_id: String,
session_configuration: SessionConfiguration,
turn_environments: TurnEnvironmentSnapshot,
options: NewTurnContextOptions,
) -> Arc<TurnContext> {
self.new_turn_context_from_configuration(
sub_id,
session_configuration,
turn_environments,
options,
TurnMultiAgentRuntime::ResolveAndStore,
self.git_enrichment_policy,
@@ -1036,9 +1046,11 @@ impl Session {
sub_id: String,
session_configuration: SessionConfiguration,
) -> Arc<TurnContext> {
let turn_environments = self.services.turn_environments.snapshot().await;
self.new_turn_context_from_configuration(
sub_id,
session_configuration,
turn_environments,
NewTurnContextOptions::default(),
TurnMultiAgentRuntime::Preview,
GitEnrichmentPolicy::Skip,
@@ -1051,11 +1063,11 @@ impl Session {
&self,
sub_id: String,
session_configuration: SessionConfiguration,
turn_environments: TurnEnvironmentSnapshot,
options: NewTurnContextOptions,
multi_agent_runtime: TurnMultiAgentRuntime,
git_enrichment_policy: GitEnrichmentPolicy,
) -> Arc<TurnContext> {
let turn_environments = self.services.turn_environments.snapshot().await;
let primary_turn_environment = turn_environments.primary();
// TODO(anp): Migrate per-turn config and legacy TurnContext cwd consumers to PathUri so
// a foreign primary environment does not fall back to the session's host cwd.
@@ -1063,7 +1075,11 @@ impl Session {
.as_ref()
.and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok())
.unwrap_or_else(|| session_configuration.cwd().clone());
let per_turn_config = self.build_per_turn_config(&session_configuration, cwd.clone());
let per_turn_config = self.build_per_turn_config(
&session_configuration,
cwd.clone(),
turn_environments.primary_workspace_roots(),
);
let network_permission_profile = primary_turn_environment
.map(TurnEnvironment::permission_profile)
.cloned()
@@ -1215,21 +1231,30 @@ impl Session {
}
}
/// Builds a context without starting work or changing the current environments.
pub(crate) async fn new_default_turn(&self) -> Arc<TurnContext> {
self.new_turn_with_default_settings(
let session_configuration = self.default_turn_configuration().await;
let turn_environments = self.services.turn_environments.snapshot().await;
self.new_turn_from_configuration(
self.next_internal_sub_id(),
session_configuration,
turn_environments,
NewTurnContextOptions::default(),
)
.await
}
/// Prepares new work on the saved environments; context-only callers keep the current ones.
pub(crate) async fn new_turn_with_default_settings(
&self,
sub_id: String,
options: NewTurnContextOptions,
) -> Arc<TurnContext> {
let session_configuration = self.default_turn_configuration().await;
self.new_turn_from_configuration(sub_id, session_configuration, options)
let turn_environments = self
.activate_turn_environments(&session_configuration)
.await;
self.new_turn_from_configuration(sub_id, session_configuration, turn_environments, options)
.await
}
@@ -645,7 +645,7 @@ async fn automatic_admission_rechecks_plan_mode_without_committing_sparse_settin
.await
.expect("explicit settings update accepts the same patch");
assert_eq!(
session.services.turn_environments.selections(),
session.configured_environment_selections().await,
proposed_environments.environments
);
assert!(session.mcp_refresh.is_pending());
+30 -3
View File
@@ -288,7 +288,12 @@ async fn executor_stop_hook_stops_after_disconnection() -> Result<()> {
async fn executor_stop_hook_rejects_mismatched_environment() -> Result<()> {
skip_if_no_network!(Ok(()));
let fixture = executor_stop_hook_fixture().await?;
let fixture = executor_hook_fixture(
["first-turn", "switch-environments", "second-turn"]
.map(completed_turn_response)
.to_vec(),
)
.await?;
let selection = fixture.attach().await?;
fixture
.test
@@ -372,6 +377,11 @@ async fn executor_stop_hook_rejects_mismatched_environment() -> Result<()> {
},
)
.await?;
fixture
.test
.submit_text_turn("apply the new executor environment")
.await?;
fixture.wait_for_hook_call().await?;
let mut mismatched_config = fixture.test.config.clone();
let mut node_repl = mismatched_config
@@ -391,7 +401,16 @@ async fn executor_stop_hook_rejects_mismatched_environment() -> Result<()> {
.codex
.refresh_mcp_config(mismatched_config)
.await;
wait_for_mcp_server(&fixture.test.codex, "node_repl").await?;
fixture
.test
.codex
.call_mcp_tool(
"node_repl",
"js",
Some(json!({"code": "1 + 1"})),
/*meta*/ None,
)
.await?;
assert_eq!(
fixture
.test
@@ -409,7 +428,15 @@ async fn executor_stop_hook_rejects_mismatched_environment() -> Result<()> {
executor.abort();
}
assert_eq!(fixture.calls().await?.len(), 1);
assert_eq!(
fixture
.calls()
.await?
.iter()
.filter(|call| call["params"]["name"] == "turn_ended")
.count(),
2
);
Ok(())
}
+145 -4
View File
@@ -9,6 +9,7 @@ use codex_core::EnvironmentConfig;
use codex_core::EnvironmentNetworkPolicy;
use codex_core::StartThreadOptions;
use codex_core::TurnInputRequest;
use codex_core::TurnInputSubmission;
use codex_core::WaitForEnvironmentToolConfig;
use codex_core::compact::SUMMARIZATION_PROMPT;
use codex_core::config::Config;
@@ -1557,10 +1558,6 @@ async fn shared_executor_keeps_ready_capability_roots_scoped_to_each_attachment(
)
.await?;
assert_eq!(
test.codex.inspect_selected_capability_roots().ready_roots,
vec![root("first-root")]
);
assert_eq!(
second
.thread
@@ -1629,6 +1626,10 @@ async fn shared_executor_keeps_ready_capability_roots_scoped_to_each_attachment(
wait_for_event(thread, |event| matches!(event, EventMsg::TurnComplete(_))).await;
}
assert_eq!(
test.codex.inspect_selected_capability_roots().ready_roots,
vec![root("first-updated-root")]
);
let requests = response_mock.requests();
let root_fragments = requests
.iter()
@@ -1981,6 +1982,146 @@ async fn pending_attachment_installs_configuration_before_waiting_turn_resumes()
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn future_pending_environment_can_finish_without_retargeting_the_active_turn() -> Result<()> {
let server = start_mock_server().await;
let mut extensions = ExtensionRegistryBuilder::new();
extensions.thread_lifecycle_contributor(Arc::new(WaitForEnvironmentTestExtension));
extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension::default()));
let mut builder = test_codex()
.with_extensions(Arc::new(extensions.build()))
.with_config(|config| {
assert!(config.features.enable(Feature::DeferredExecutor).is_ok());
});
let test = builder.build_with_auto_env(&server).await?;
let selection = test.codex.environment_selections().await.remove(0);
let active = TurnEnvironmentSelection {
config: EnvironmentConfigState::Pending,
..selection
};
let future_cwd = active.cwd.join("future-environment")?;
test.fs()
.create_directory(
&future_cwd,
CreateDirectoryOptions {
recursive: false,
follow_symlinks: true,
},
/*sandbox*/ None,
)
.await?;
let future = TurnEnvironmentSelection {
cwd: future_cwd.clone(),
workspace_roots: vec![future_cwd],
..active.clone()
};
let owner_config = |selection: &TurnEnvironmentSelection, id: &str| EnvironmentConfig {
allow_login_shell: test.config.permissions.allow_login_shell,
workspace_roots: selection.workspace_roots.clone(),
permission_profile: PermissionProfileSnapshot::legacy(PermissionProfile::read_only()),
shell_environment_policy: Default::default(),
windows_sandbox_level: WindowsSandboxLevel::from_config(&test.config),
windows_sandbox_private_desktop: test.config.permissions.windows_sandbox_private_desktop,
use_legacy_landlock: test.config.features.use_legacy_landlock(),
exec_policy: None,
mcp_policy: None,
network_policy: None,
selected_capability_roots: vec![SelectedCapabilityRoot {
id: id.to_string(),
location: CapabilityRootLocation::Environment {
environment_id: selection.environment_id.clone(),
path: selection.cwd.clone(),
},
}],
};
let thread = test
.thread_manager
.start_thread(StartThreadOptions {
environments: Some(vec![active.clone()]),
..StartThreadOptions::new(test.config.clone())
})
.await?
.thread;
let response_mock = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_function_call(
"wait-for-active-environment",
"wait_for_environment",
&json!({ "environment_id": active.environment_id }).to_string(),
),
ev_completed("waiting"),
]),
sse(vec![ev_completed("active-done")]),
sse(vec![ev_completed("future-done")]),
],
)
.await;
let request = |text: &str| {
TurnInputRequest::user_input(vec![UserInput::Text {
text: text.to_string(),
text_elements: Vec::new(),
}])
};
thread.start_or_steer_turn(request("wait")).await?;
wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await;
assert!(matches!(
thread
.start_or_steer_turn(
request("use the other environment next turn").with_thread_settings(
ThreadSettingsOverrides {
environments: Some(TurnEnvironmentSelections::new(
test.config.cwd.clone(),
vec![future.clone()],
)),
..Default::default()
},
)
)
.await?,
TurnInputSubmission::Steered { .. }
));
thread
.environment_ready(&future, owner_config(&future, "future-root"))
.await?;
assert!(matches!(
thread.environment_selections().await[0].config,
EnvironmentConfigState::Ready(_)
));
thread
.environment_ready(&active, owner_config(&active, "active-root"))
.await?;
wait_for_event(&thread, |event| matches!(event, EventMsg::TurnComplete(_))).await;
thread.start_or_steer_turn(request("next turn")).await?;
wait_for_event(&thread, |event| matches!(event, EventMsg::TurnComplete(_))).await;
let requests = response_mock.requests();
let (output, _) = requests[1]
.function_call_output_content_and_success("wait-for-active-environment")
.context("active wait output should be model visible")?;
assert_eq!(
serde_json::from_str::<Value>(&output.context("active wait output should contain JSON")?)?,
json!({ "environment_id": active.environment_id, "status": "ready" })
);
let ready_roots = requests[1..]
.iter()
.map(|request| {
request
.message_input_texts("user")
.into_iter()
.rfind(|text| text.contains("<ready_capability_roots>"))
})
.collect::<Vec<_>>();
assert_eq!(
ready_roots,
["active-root", "future-root"].map(|root| Some(format!(
"<ready_capability_roots>{root}</ready_capability_roots>"
)))
);
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)]
+13 -1
View File
@@ -704,7 +704,7 @@ async fn environment_mcp_policy_filters_runtime_config_and_model_tools(
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let response = mount_sse_once(
mount_sse_once(
&server,
responses::sse(vec![
responses::ev_response_created("resp-1"),
@@ -786,6 +786,9 @@ async fn environment_mcp_policy_filters_runtime_config_and_model_tools(
},
)
.await?;
fixture
.submit_text_turn("start with pending environment configuration")
.await?;
let (pending_config, _) = fixture.codex.current_mcp_config_and_runtime_context().await;
let pending_servers = pending_config.mcp_server_catalog.configured_servers();
assert!(!pending_servers["allowed"].enabled);
@@ -856,6 +859,15 @@ async fn environment_mcp_policy_filters_runtime_config_and_model_tools(
)
.await?;
let response = mount_sse_once(
&server,
responses::sse(vec![
responses::ev_response_created("resp-2"),
responses::ev_assistant_message("msg-2", "done"),
responses::ev_completed("resp-2"),
]),
)
.await;
fixture
.submit_text_turn("show the available MCP tools")
.await?;