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

Mark the primary environment in model context (#35874)

Browse Source 
## Why

When a turn uses multiple environments, the environment context lists each one but did not identify which environment is primary.

## What changed

- Add a `primary` attribute to each environment in multi-environment model context.
- Persist primary selection in world-state snapshots and emit updates when the primary environment changes.
- Restate environments when transitioning between single- and multi-environment context, while preserving the existing single-environment format.

## Testing

- Cover primary-environment changes, single/multi-environment transitions, legacy snapshot compatibility, and deferred-executor promotion after environment startup.

GitOrigin-RevId: 73a1d794a227edef6e9b7844e7f947b7f598dc80
This commit is contained in:
sayan-oai
2026-07-29 03:02:32 +00:00
committed by copyberry
parent 9f23e97797
commit 0a6616f4cf
5 changed files with 370 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/context/world_state/environment.rs
+25 -5
View File
@@ -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 = &current.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::<BTreeMap<_, _>>();
@@ -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<String, EnvironmentUpdate>,
legacy_single: bool,
include_primary: bool,
current_date: Option<String>,
timezone: Option<String>,
network: Option<NetworkContext>,
@@ -210,6 +216,13 @@ impl ContextualUserFragment for RenderedEnvironments {
rendered.push_str(" <environment id=\"");
push_xml_escaped_text(&mut rendered, id);
rendered.push('"');
if self.include_primary {
rendered.push_str(if environment.is_primary {
" primary=\"true\""
} else {
" primary=\"false\""
});
}
rendered.push_str(">\n");
push_environment_values(&mut rendered, environment, " ");
rendered.push_str(" </environment>\n");
@@ -283,6 +296,7 @@ struct EnvironmentState {
cwd: PathUri,
status: EnvironmentStatus,
shell: Option<String>,
is_primary: bool,
}
#[derive(Default, Deserialize, Serialize)]
@@ -300,12 +314,15 @@ struct EnvironmentSnapshot {
cwd: String,
status: EnvironmentStatus,
shell: Option<String>,
#[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<String, EnvironmentState> {
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<String, En
.shell
.as_ref()
.map(|shell| shell.name().to_string()),
is_primary: index == 0,
},
)
})
@@ -345,6 +364,7 @@ fn environment_states(snapshot: &TurnEnvironmentSnapshot) -> BTreeMap<String, En
cwd: environment.selection.cwd.clone(),
status: EnvironmentStatus::Starting,
shell: None,
is_primary: false,
});
}
environments
codex-rs/core/src/context/world_state/environment_render_tests.rs
+14 -5
View File
@@ -35,6 +35,7 @@ fn environment(id: &str, cwd: PathUri, shell: impl Into<String>) -> (String, Env
cwd,
status: EnvironmentStatus::Available,
shell: Some(shell.into()),
is_primary: false,
},
)
}
@@ -46,8 +47,16 @@ fn environment_state(
network: Option<NetworkContext>,
subagents: Option<String>,
) -> 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#"<environment_context>
<environments>
<environment id="local">
<environment id="local" primary="true">
<cwd>{}</cwd>
<shell>bash</shell>
</environment>
<environment id="remote">
<environment id="remote" primary="false">
<cwd>{}</cwd>
<shell>bash</shell>
</environment>
@@ -328,11 +337,11 @@ fn serialize_environment_context_prefers_environment_shell_when_present() {
let expected = format!(
r#"<environment_context>
<environments>
<environment id="local">
<environment id="local" primary="true">
<cwd>{}</cwd>
<shell>powershell</shell>
</environment>
<environment id="remote">
<environment id="remote" primary="false">
<cwd>{}</cwd>
<shell>cmd</shell>
</environment>
codex-rs/core/src/context/world_state/environment_tests.rs
+166 -3
View File
@@ -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!(
"<environment_context>\n <environments>\n <environment id=\"local\" primary=\"false\">\n <cwd>{}</cwd>\n <shell>bash</shell>\n </environment>\n <environment id=\"remote\" primary=\"true\">\n <cwd>{}</cwd>\n <shell>zsh</shell>\n </environment>\n </environments>\n</environment_context>",
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::<EnvironmentsSnapshot>(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!(
"<environment_context>\n <environments>\n <environment id=\"local\" primary=\"{local_is_primary}\">\n <cwd>{local_cwd}</cwd>\n <shell>bash</shell>\n </environment>\n <environment id=\"remote\" primary=\"{remote_is_primary}\">\n <cwd>{remote_cwd}</cwd>\n <shell>zsh</shell>\n </environment>\n </environments>\n</environment_context>"
)
);
let reduced = single
.render_diff(PreviousSectionState::Known(&multiple.snapshot()))
.expect("removing an environment should update the model")
.render();
assert_eq!(
reduced,
format!(
"<environment_context>\n <environments>\n <environment id=\"local\" primary=\"true\">\n <cwd>{local_cwd}</cwd>\n <shell>bash</shell>\n </environment>\n <environment id=\"remote\" status=\"unavailable\" />\n </environments>\n</environment_context>"
)
);
}
Ok(())
}
fn available(cwd: &str, shell: &str) -> Result<EnvironmentState> {
Ok(EnvironmentState {
cwd: PathUri::parse(cwd)?,
status: EnvironmentStatus::Available,
shell: Some(shell.to_string()),
is_primary: false,
})
}
fn primary(cwd: &str, shell: &str) -> Result<EnvironmentState> {
Ok(EnvironmentState {
is_primary: true,
..available(cwd, shell)?
})
}
@@ -162,5 +324,6 @@ fn starting(cwd: &str) -> Result<EnvironmentState> {
cwd: PathUri::parse(cwd)?,
status: EnvironmentStatus::Starting,
shell: None,
is_primary: false,
})
}
codex-rs/core/src/context/world_state/snapshots/codex_core__context__world_state__environment__tests__snapshots.snap
+10 -10
View File
@@ -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)
<environment_context>
<environments>
<environment id="devbox">
<environment id="devbox" primary="false">
<cwd>/workspace</cwd>
<shell>bash</shell>
</environment>
<environment id="laptop">
<environment id="laptop" primary="true">
<cwd>/repo</cwd>
<shell>zsh</shell>
</environment>
</environments>
</environment_context>
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)
<environment_context>
<environments>
<environment id="devbox">
<environment id="devbox" primary="false">
<cwd>/workspace</cwd>
<shell>bash</shell>
</environment>
<environment id="laptop">
<environment id="laptop" primary="true">
<cwd>/repo</cwd>
<shell>zsh</shell>
</environment>
</environments>
</environment_context>
{"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)
<environment_context>
<environments>
<environment id="devbox">
<environment id="devbox" primary="false">
<cwd>/workspace</cwd>
<shell>powershell</shell>
</environment>
<environment id="laptop">
<environment id="laptop" primary="true">
<cwd>/repo</cwd>
<shell>zsh</shell>
</environment>
<environment id="old" status="unavailable" />
<environment id="remote">
<environment id="remote" primary="false">
<cwd>/remote</cwd>
<status>starting</status>
</environment>
codex-rs/core/tests/suite/remote_env.rs
+155
View File
@@ -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("<environment_context>"))
.context("starting environment context")?;
assert!(initial_context.contains("<environment id=\"local\" primary=\"true\">"));
assert!(initial_context.contains("<environment id=\"remote\" primary=\"false\">"));
assert!(initial_context.contains("<status>starting</status>"));
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("<environment_context>"))
.context("updated primary environment context")?;
assert!(updated_context.contains("<environment id=\"local\" primary=\"false\">"));
assert!(updated_context.contains("<environment id=\"remote\" primary=\"true\">"));
assert!(updated_context.contains("<shell>zsh</shell>"));
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::<RolloutLine>)
.collect::<serde_json::Result<Vec<_>>>()?
.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<TcpStream>) -> Value {
loop {
match timeout(Duration::from_secs(5), websocket.next())