diff --git a/codex-rs/core/src/context/world_state/environment.rs b/codex-rs/core/src/context/world_state/environment.rs index d5a8abf78d..735f864a6c 100644 --- a/codex-rs/core/src/context/world_state/environment.rs +++ b/codex-rs/core/src/context/world_state/environment.rs @@ -62,6 +62,7 @@ impl EnvironmentsState { }) .collect(), legacy_single: is_legacy_single(&self.environments), + include_primary: self.environments.len() > 1, current_date: self.current_date.clone(), timezone: self.timezone.clone(), network: self.network.clone(), @@ -87,6 +88,7 @@ impl WorldStateSection for EnvironmentsState { cwd: environment.cwd.inferred_native_path_string(), status: environment.status, shell: environment.shell.clone(), + is_primary: self.environments.len() > 1 && environment.is_primary, }, ) }) @@ -113,15 +115,17 @@ impl WorldStateSection for EnvironmentsState { || current.timezone != previous.timezone || current.network != previous.network || current.filesystem != previous.filesystem; + let multiple_environments = self.environments.len() > 1; + let previous_multiple_environments = previous.environments.len() > 1; let mut updates = self .environments .iter() .filter(|(id, _)| { let environment = ¤t.environments[*id]; - previous - .environments - .get(*id) - .is_none_or(|previous| !environment.has_same_diff_value(previous)) + previous.environments.get(*id).is_none_or(|previous| { + multiple_environments != previous_multiple_environments + || !environment.has_same_diff_value(previous) + }) }) .map(|(id, environment)| (id.clone(), EnvironmentUpdate::Current(environment.clone()))) .collect::>(); @@ -140,6 +144,7 @@ impl WorldStateSection for EnvironmentsState { Box::new(RenderedEnvironments { updates, legacy_single, + include_primary: multiple_environments || previous_multiple_environments, current_date: self.current_date.clone(), timezone: self.timezone.clone(), network: self.network.clone(), @@ -171,6 +176,7 @@ impl ContextualUserFragment for EnvironmentsState { struct RenderedEnvironments { updates: BTreeMap, legacy_single: bool, + include_primary: bool, current_date: Option, timezone: Option, network: Option, @@ -210,6 +216,13 @@ impl ContextualUserFragment for RenderedEnvironments { rendered.push_str(" \n"); push_environment_values(&mut rendered, environment, " "); rendered.push_str(" \n"); @@ -283,6 +296,7 @@ struct EnvironmentState { cwd: PathUri, status: EnvironmentStatus, shell: Option, + is_primary: bool, } #[derive(Default, Deserialize, Serialize)] @@ -300,12 +314,15 @@ struct EnvironmentSnapshot { cwd: String, status: EnvironmentStatus, shell: Option, + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + is_primary: bool, } impl EnvironmentSnapshot { fn has_same_diff_value(&self, other: &Self) -> bool { self.cwd == other.cwd && self.status == other.status + && self.is_primary == other.is_primary && self .shell .as_ref() @@ -324,7 +341,8 @@ enum EnvironmentStatus { fn environment_states(snapshot: &TurnEnvironmentSnapshot) -> BTreeMap { let mut environments = snapshot .turn_environments() - .map(|environment| { + .enumerate() + .map(|(index, environment)| { ( environment.environment_id.clone(), EnvironmentState { @@ -334,6 +352,7 @@ fn environment_states(snapshot: &TurnEnvironmentSnapshot) -> BTreeMap BTreeMap) -> (String, Env cwd, status: EnvironmentStatus::Available, shell: Some(shell.into()), + is_primary: false, }, ) } @@ -46,8 +47,16 @@ fn environment_state( network: Option, subagents: Option, ) -> EnvironmentsState { + let environments = environments + .into_iter() + .enumerate() + .map(|(index, (id, mut environment))| { + environment.is_primary = index == 0; + (id, environment) + }) + .collect(); EnvironmentsState { - environments: environments.into_iter().collect(), + environments, current_date, timezone, network, @@ -287,11 +296,11 @@ fn serialize_environment_context_with_multiple_selected_environments() { let expected = format!( r#" - + {} bash - + {} bash @@ -328,11 +337,11 @@ fn serialize_environment_context_prefers_environment_shell_when_present() { let expected = format!( r#" - + {} powershell - + {} cmd diff --git a/codex-rs/core/src/context/world_state/environment_tests.rs b/codex-rs/core/src/context/world_state/environment_tests.rs index 0e666ec64f..9a9c66e483 100644 --- a/codex-rs/core/src/context/world_state/environment_tests.rs +++ b/codex-rs/core/src/context/world_state/environment_tests.rs @@ -5,6 +5,8 @@ use anyhow::Result; use codex_exec_server::LOCAL_ENVIRONMENT_ID; use codex_protocol::models::PermissionProfile; use codex_protocol::permissions::NetworkSandboxPolicy; +use pretty_assertions::assert_eq; +use serde_json::json; #[test] fn snapshots() -> Result<()> { @@ -14,7 +16,7 @@ fn snapshots() -> Result<()> { let full = EnvironmentsState { environments: [ - ("laptop".to_string(), available("file:///repo", "zsh")?), + ("laptop".to_string(), primary("file:///repo", "zsh")?), ( "devbox".to_string(), available("file:///workspace", "bash")?, @@ -26,7 +28,7 @@ fn snapshots() -> Result<()> { }; let before_environment_changes = EnvironmentsState { environments: [ - ("laptop".to_string(), available("file:///repo", "bash")?), + ("laptop".to_string(), primary("file:///repo", "bash")?), ("devbox".to_string(), starting("file:///workspace")?), ("old".to_string(), available("file:///old", "sh")?), ] @@ -36,7 +38,7 @@ fn snapshots() -> Result<()> { }; let after_environment_changes = EnvironmentsState { environments: [ - ("laptop".to_string(), available("file:///repo", "zsh")?), + ("laptop".to_string(), primary("file:///repo", "zsh")?), ( "devbox".to_string(), available("file:///workspace", "powershell")?, @@ -104,6 +106,7 @@ fn snapshots() -> Result<()> { cwd: PathUri::parse("file:///repo")?, status: EnvironmentStatus::Available, shell: None, + is_primary: false, }, )] .into_iter() @@ -149,11 +152,170 @@ fn snapshots() -> Result<()> { Ok(()) } +#[test] +fn changing_primary_environment_updates_model_context_and_persisted_state() -> Result<()> { + let before = EnvironmentsState { + environments: [ + ("local".to_string(), primary("file:///local", "bash")?), + ("remote".to_string(), available("file:///remote", "zsh")?), + ] + .into_iter() + .collect(), + ..Default::default() + }; + let after = EnvironmentsState { + environments: [ + ("local".to_string(), available("file:///local", "bash")?), + ("remote".to_string(), primary("file:///remote", "zsh")?), + ] + .into_iter() + .collect(), + ..Default::default() + }; + let previous = before.snapshot(); + let rendered = after + .render_diff(PreviousSectionState::Known(&previous)) + .expect("primary change should update the model") + .render(); + + assert_eq!( + rendered, + format!( + "\n \n \n {}\n bash\n \n \n {}\n zsh\n \n \n", + PathUri::parse("file:///local")?.inferred_native_path_string(), + PathUri::parse("file:///remote")?.inferred_native_path_string(), + ) + ); + + let mut previous_world_state = super::super::WorldState::default(); + previous_world_state.add_section(before); + let mut current_world_state = super::super::WorldState::default(); + current_world_state.add_section(after); + assert_eq!( + current_world_state + .snapshot() + .merge_patch_from(&previous_world_state.snapshot()), + Some(json!({ + "environments": { + "environments": { + "local": { "is_primary": null }, + "remote": { "is_primary": true }, + }, + }, + })) + ); + + Ok(()) +} + +#[test] +fn legacy_single_environment_snapshot_does_not_change() -> Result<()> { + let environment = EnvironmentsState { + environments: [("local".to_string(), primary("file:///repo", "bash")?)] + .into_iter() + .collect(), + ..Default::default() + }; + let legacy_snapshot = serde_json::from_value::(json!({ + "environments": { + "local": { + "cwd": PathUri::parse("file:///repo")?.inferred_native_path_string(), + "status": "available", + "shell": "bash", + }, + }, + }))?; + + assert!( + environment + .render_diff(PreviousSectionState::Known(&legacy_snapshot)) + .is_none() + ); + assert_eq!( + serde_json::to_value(environment.snapshot())?["environments"]["local"], + json!({ + "cwd": PathUri::parse("file:///repo")?.inferred_native_path_string(), + "status": "available", + "shell": "bash", + }) + ); + + Ok(()) +} + +#[test] +fn crossing_single_environment_boundary_restates_current_environments() -> Result<()> { + let single = EnvironmentsState { + environments: [("local".to_string(), primary("file:///local", "bash")?)] + .into_iter() + .collect(), + ..Default::default() + }; + let local_cwd = PathUri::parse("file:///local")?.inferred_native_path_string(); + let remote_cwd = PathUri::parse("file:///remote")?.inferred_native_path_string(); + + for (local_is_primary, remote_is_primary) in [(true, false), (false, true)] { + let multiple = EnvironmentsState { + environments: [ + ( + "local".to_string(), + EnvironmentState { + is_primary: local_is_primary, + ..available("file:///local", "bash")? + }, + ), + ( + "remote".to_string(), + EnvironmentState { + is_primary: remote_is_primary, + ..available("file:///remote", "zsh")? + }, + ), + ] + .into_iter() + .collect(), + ..Default::default() + }; + + let expanded = multiple + .render_diff(PreviousSectionState::Known(&single.snapshot())) + .expect("adding an environment should update the model") + .render(); + assert_eq!( + expanded, + format!( + "\n \n \n {local_cwd}\n bash\n \n \n {remote_cwd}\n zsh\n \n \n" + ) + ); + + let reduced = single + .render_diff(PreviousSectionState::Known(&multiple.snapshot())) + .expect("removing an environment should update the model") + .render(); + assert_eq!( + reduced, + format!( + "\n \n \n {local_cwd}\n bash\n \n \n \n" + ) + ); + } + + Ok(()) +} + fn available(cwd: &str, shell: &str) -> Result { Ok(EnvironmentState { cwd: PathUri::parse(cwd)?, status: EnvironmentStatus::Available, shell: Some(shell.to_string()), + is_primary: false, + }) +} + +fn primary(cwd: &str, shell: &str) -> Result { + Ok(EnvironmentState { + is_primary: true, + ..available(cwd, shell)? }) } @@ -162,5 +324,6 @@ fn starting(cwd: &str) -> Result { cwd: PathUri::parse(cwd)?, status: EnvironmentStatus::Starting, shell: None, + is_primary: false, }) } diff --git a/codex-rs/core/src/context/world_state/snapshots/codex_core__context__world_state__environment__tests__snapshots.snap b/codex-rs/core/src/context/world_state/snapshots/codex_core__context__world_state__environment__tests__snapshots.snap index 60e5ba315c..cce3528bcf 100644 --- a/codex-rs/core/src/context/world_state/snapshots/codex_core__context__world_state__environment__tests__snapshots.snap +++ b/codex-rs/core/src/context/world_state/snapshots/codex_core__context__world_state__environment__tests__snapshots.snap @@ -5,47 +5,47 @@ expression: "render_section_cases(&[(Absent, Absent), (Absent, Known(&full)),\n( Absent -> Absent None -Absent -> {"environments":{"devbox":{"cwd":"/workspace","shell":"bash","status":"available"},"laptop":{"cwd":"/repo","shell":"zsh","status":"available"}}} (role - user) +Absent -> {"environments":{"devbox":{"cwd":"/workspace","shell":"bash","status":"available"},"laptop":{"cwd":"/repo","is_primary":true,"shell":"zsh","status":"available"}}} (role - user) - + /workspace bash - + /repo zsh -Unknown -> {"environments":{"devbox":{"cwd":"/workspace","shell":"bash","status":"available"},"laptop":{"cwd":"/repo","shell":"zsh","status":"available"}}} (role - user) +Unknown -> {"environments":{"devbox":{"cwd":"/workspace","shell":"bash","status":"available"},"laptop":{"cwd":"/repo","is_primary":true,"shell":"zsh","status":"available"}}} (role - user) - + /workspace bash - + /repo zsh -{"environments":{"devbox":{"cwd":"/workspace","status":"starting"},"laptop":{"cwd":"/repo","shell":"bash","status":"available"},"old":{"cwd":"/old","shell":"sh","status":"available"}}} -> {"environments":{"devbox":{"cwd":"/workspace","shell":"powershell","status":"available"},"laptop":{"cwd":"/repo","shell":"zsh","status":"available"},"remote":{"cwd":"/remote","status":"starting"}}} (role - user) +{"environments":{"devbox":{"cwd":"/workspace","status":"starting"},"laptop":{"cwd":"/repo","is_primary":true,"shell":"bash","status":"available"},"old":{"cwd":"/old","shell":"sh","status":"available"}}} -> {"environments":{"devbox":{"cwd":"/workspace","shell":"powershell","status":"available"},"laptop":{"cwd":"/repo","is_primary":true,"shell":"zsh","status":"available"},"remote":{"cwd":"/remote","status":"starting"}}} (role - user) - + /workspace powershell - + /repo zsh - + /remote starting diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index ba07c6ca9e..4b8680e56f 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -548,6 +548,161 @@ async fn settings_update_does_not_retarget_active_turn_environment() -> Result<( Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn deferred_executor_promotes_primary_environment_when_startup_completes() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let server = start_mock_server().await; + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("warmup"), + ev_assistant_message("warmup-message", "ready"), + ev_completed("warmup"), + ]), + sse(vec![ + ev_response_created("before-promotion"), + ev_function_call( + "pause-for-environment", + "request_user_input", + &json!({ + "questions": [{ + "id": "continue", + "header": "Continue", + "question": "Continue after the environment starts?", + "options": [{ + "label": "Yes (Recommended)", + "description": "Continue the test." + }, { + "label": "No", + "description": "Stop the test." + }] + }] + }) + .to_string(), + ), + ev_completed("before-promotion"), + ]), + sse(vec![ + ev_response_created("after-promotion"), + ev_assistant_message("after-promotion-message", "done"), + ev_completed("after-promotion"), + ]), + ], + ) + .await; + let mut builder = test_codex() + .with_exec_server_url(format!("ws://{}", listener.local_addr()?)) + .with_config(|config| { + config.project_doc_max_bytes = 0; + assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); + assert!( + config + .features + .enable(Feature::DefaultModeRequestUserInput) + .is_ok() + ); + }); + let test = builder.build_with_remote_and_local_env(&server).await?; + let local_selection = local(test.config.cwd.clone()); + let remote_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)], + }; + + test.submit_turn_with_environments( + "warm the local environment", + Some(vec![local_selection.clone(), remote_selection.clone()]), + ) + .await?; + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "wait for the primary environment".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: ThreadSettingsOverrides { + environments: Some(TurnEnvironmentSelections::new( + test.config.cwd.clone(), + vec![remote_selection, local_selection], + )), + ..Default::default() + }, + }) + .await?; + let request = wait_for_event_match(&test.codex, |event| match event { + EventMsg::RequestUserInput(request) => Some(request.clone()), + _ => None, + }) + .await; + + let requests = response_mock.requests(); + let initial_context = requests[1] + .message_input_texts("user") + .into_iter() + .rfind(|text| text.contains("")) + .context("starting environment context")?; + assert!(initial_context.contains("")); + assert!(initial_context.contains("")); + assert!(initial_context.contains("starting")); + + serve_environment_info(listener).await; + test.codex + .submit(Op::UserInputAnswer { + id: request.turn_id, + response: RequestUserInputResponse { + answers: HashMap::from([( + "continue".to_string(), + RequestUserInputAnswer { + answers: vec!["Yes (Recommended)".to_string()], + }, + )]), + }, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = response_mock.requests(); + let updated_context = requests[2] + .message_input_texts("user") + .into_iter() + .rfind(|text| text.contains("")) + .context("updated primary environment context")?; + assert!(updated_context.contains("")); + assert!(updated_context.contains("")); + assert!(updated_context.contains("zsh")); + + test.codex.ensure_rollout_materialized().await; + test.codex.flush_rollout().await?; + let rollout = fs::read_to_string(test.codex.rollout_path().context("rollout path")?)?; + let world_state_patch = rollout + .lines() + .map(serde_json::from_str::) + .collect::>>()? + .into_iter() + .filter_map(|line| match line.item { + RolloutItem::WorldState(item) if !item.full => Some(item.state), + _ => None, + }) + .find(|patch| { + patch.pointer("/environments/environments/remote/is_primary") == Some(&json!(true)) + }) + .context("primary environment World State patch")?; + assert_eq!( + world_state_patch.pointer("/environments/environments/local/is_primary"), + Some(&Value::Null) + ); + + Ok(()) +} + async fn read_exec_server_json(websocket: &mut WebSocketStream) -> Value { loop { match timeout(Duration::from_secs(5), websocket.next())