[prompt] xml-format EnvironmentContext (#2272)

## Summary
Before we land #2243, let's start printing environment_context in our
preferred format. This struct will evolve over time with new
information, xml gives us a balance of human readable without too much
parsing, llm readable, and extensible.

Also moves us over to an Option-based struct, so we can easily provide
diffs to the model.

## Testing
- [x] Updated tests to reflect new format

Author: dylan-hurd-oai

codex-rs/core/src/codex.rs

@@ -508,10 +508,10 @@ impl Session {
             conversation_items.push(Prompt::format_user_instructions_message(user_instructions));
         }
         conversation_items.push(ResponseItem::from(EnvironmentContext::new(
-            turn_context.cwd.to_path_buf(),
-            turn_context.approval_policy,
-            turn_context.sandbox_policy.clone(),
-            sess.user_shell.clone(),
+            Some(turn_context.cwd.clone()),
+            Some(turn_context.approval_policy),
+            Some(turn_context.sandbox_policy.clone()),
+            Some(sess.user_shell.clone()),
         )));
         sess.record_conversation_items(&conversation_items).await;
 
@@ -1068,10 +1068,11 @@ async fn submission_loop(
                 turn_context = Arc::new(new_turn_context);
                 if cwd.is_some() || approval_policy.is_some() || sandbox_policy.is_some() {
                     sess.record_conversation_items(&[ResponseItem::from(EnvironmentContext::new(
-                        new_cwd,
-                        new_approval_policy,
-                        new_sandbox_policy,
-                        sess.user_shell.clone(),
+                        cwd,
+                        approval_policy,
+                        sandbox_policy,
+                        // Shell is not configurable from turn to turn
+                        None,
                     ))])
                     .await;
                 }

codex-rs/core/src/environment_context.rs

@@ -8,11 +8,10 @@ use crate::protocol::AskForApproval;
 use crate::protocol::SandboxPolicy;
 use crate::shell::Shell;
 use codex_protocol::config_types::SandboxMode;
-use std::fmt::Display;
 use std::path::PathBuf;
 
 /// wraps environment context message in a tag for the model to parse more easily.
-pub(crate) const ENVIRONMENT_CONTEXT_START: &str = "<environment_context>\n";
+pub(crate) const ENVIRONMENT_CONTEXT_START: &str = "<environment_context>";
 pub(crate) const ENVIRONMENT_CONTEXT_END: &str = "</environment_context>";
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, DeriveDisplay)]
@@ -25,58 +24,87 @@ pub enum NetworkAccess {
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 #[serde(rename = "environment_context", rename_all = "snake_case")]
 pub(crate) struct EnvironmentContext {
-    pub cwd: PathBuf,
-    pub approval_policy: AskForApproval,
-    pub sandbox_mode: SandboxMode,
-    pub network_access: NetworkAccess,
-    pub shell: Shell,
+    pub cwd: Option<PathBuf>,
+    pub approval_policy: Option<AskForApproval>,
+    pub sandbox_mode: Option<SandboxMode>,
+    pub network_access: Option<NetworkAccess>,
+    pub shell: Option<Shell>,
 }
 
 impl EnvironmentContext {
     pub fn new(
-        cwd: PathBuf,
-        approval_policy: AskForApproval,
-        sandbox_policy: SandboxPolicy,
-        shell: Shell,
+        cwd: Option<PathBuf>,
+        approval_policy: Option<AskForApproval>,
+        sandbox_policy: Option<SandboxPolicy>,
+        shell: Option<Shell>,
     ) -> Self {
         Self {
             cwd,
             approval_policy,
             sandbox_mode: match sandbox_policy {
-                SandboxPolicy::DangerFullAccess => SandboxMode::DangerFullAccess,
-                SandboxPolicy::ReadOnly => SandboxMode::ReadOnly,
-                SandboxPolicy::WorkspaceWrite { .. } => SandboxMode::WorkspaceWrite,
+                Some(SandboxPolicy::DangerFullAccess) => Some(SandboxMode::DangerFullAccess),
+                Some(SandboxPolicy::ReadOnly) => Some(SandboxMode::ReadOnly),
+                Some(SandboxPolicy::WorkspaceWrite { .. }) => Some(SandboxMode::WorkspaceWrite),
+                None => None,
             },
             network_access: match sandbox_policy {
-                SandboxPolicy::DangerFullAccess => NetworkAccess::Enabled,
-                SandboxPolicy::ReadOnly => NetworkAccess::Restricted,
-                SandboxPolicy::WorkspaceWrite { network_access, .. } => {
+                Some(SandboxPolicy::DangerFullAccess) => Some(NetworkAccess::Enabled),
+                Some(SandboxPolicy::ReadOnly) => Some(NetworkAccess::Restricted),
+                Some(SandboxPolicy::WorkspaceWrite { network_access, .. }) => {
                     if network_access {
-                        NetworkAccess::Enabled
+                        Some(NetworkAccess::Enabled)
                     } else {
-                        NetworkAccess::Restricted
+                        Some(NetworkAccess::Restricted)
                     }
                 }
+                None => None,
             },
             shell,
         }
     }
 }
 
-impl Display for EnvironmentContext {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        writeln!(
-            f,
-            "Current working directory: {}",
-            self.cwd.to_string_lossy()
-        )?;
-        writeln!(f, "Approval policy: {}", self.approval_policy)?;
-        writeln!(f, "Sandbox mode: {}", self.sandbox_mode)?;
-        writeln!(f, "Network access: {}", self.network_access)?;
-        if let Some(shell_name) = self.shell.name() {
-            writeln!(f, "Shell: {shell_name}")?;
+impl EnvironmentContext {
+    /// Serializes the environment context to XML. Libraries like `quick-xml`
+    /// require custom macros to handle Enums with newtypes, so we just do it
+    /// manually, to keep things simple. Output looks like:
+    ///
+    /// ```xml
+    /// <environment_context>
+    ///   <cwd>...</cwd>
+    ///   <approval_policy>...</approval_policy>
+    ///   <sandbox_mode>...</sandbox_mode>
+    ///   <network_access>...</network_access>
+    ///   <shell>...</shell>
+    /// </environment_context>
+    /// ```
+    pub fn serialize_to_xml(self) -> String {
+        let mut lines = vec![ENVIRONMENT_CONTEXT_START.to_string()];
+        if let Some(cwd) = self.cwd {
+            lines.push(format!("  <cwd>{}</cwd>", cwd.to_string_lossy()));
+        }
+        if let Some(approval_policy) = self.approval_policy {
+            lines.push(format!(
+                "  <approval_policy>{}</approval_policy>",
+                approval_policy
+            ));
+        }
+        if let Some(sandbox_mode) = self.sandbox_mode {
+            lines.push(format!("  <sandbox_mode>{}</sandbox_mode>", sandbox_mode));
+        }
+        if let Some(network_access) = self.network_access {
+            lines.push(format!(
+                "  <network_access>{}</network_access>",
+                network_access
+            ));
+        }
+        if let Some(shell) = self.shell
+            && let Some(shell_name) = shell.name()
+        {
+            lines.push(format!("  <shell>{}</shell>", shell_name));
         }
-        Ok(())
+        lines.push(ENVIRONMENT_CONTEXT_END.to_string());
+        lines.join("\n")
     }
 }
 
@@ -86,7 +114,7 @@ impl From<EnvironmentContext> for ResponseItem {
             id: None,
             role: "user".to_string(),
             content: vec![ContentItem::InputText {
-                text: format!("{ENVIRONMENT_CONTEXT_START}{ec}{ENVIRONMENT_CONTEXT_END}"),
+                text: ec.serialize_to_xml(),
             }],
         }
     }

codex-rs/core/tests/prompt_caching.rs

@@ -89,10 +89,15 @@ async fn prefixes_context_and_instructions_once_and_consistently_across_requests
     let shell = default_user_shell().await;
 
     let expected_env_text = format!(
-        "<environment_context>\nCurrent working directory: {}\nApproval policy: on-request\nSandbox mode: read-only\nNetwork access: restricted\n{}</environment_context>",
+        r#"<environment_context>
+  <cwd>{}</cwd>
+  <approval_policy>on-request</approval_policy>
+  <sandbox_mode>read-only</sandbox_mode>
+  <network_access>restricted</network_access>
+{}</environment_context>"#,
         cwd.path().to_string_lossy(),
         match shell.name() {
-            Some(name) => format!("Shell: {name}\n"),
+            Some(name) => format!("  <shell>{}</shell>\n", name),
             None => String::new(),
         }
     );
@@ -190,12 +195,10 @@ async fn overrides_turn_context_but_keeps_cached_prefix_and_key_constant() {
         .unwrap();
     wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
 
-    // Change everything about the turn context.
-    let new_cwd = TempDir::new().unwrap();
     let writable = TempDir::new().unwrap();
     codex
         .submit(Op::OverrideTurnContext {
-            cwd: Some(new_cwd.path().to_path_buf()),
+            cwd: None,
             approval_policy: Some(AskForApproval::Never),
             sandbox_policy: Some(SandboxPolicy::WorkspaceWrite {
                 writable_roots: vec![writable.path().to_path_buf()],
@@ -227,7 +230,6 @@ async fn overrides_turn_context_but_keeps_cached_prefix_and_key_constant() {
 
     let body1 = requests[0].body_json::<serde_json::Value>().unwrap();
     let body2 = requests[1].body_json::<serde_json::Value>().unwrap();
-
     // prompt_cache_key should remain constant across overrides
     assert_eq!(
         body1["prompt_cache_key"], body2["prompt_cache_key"],
@@ -243,16 +245,13 @@ async fn overrides_turn_context_but_keeps_cached_prefix_and_key_constant() {
         "content": [ { "type": "input_text", "text": "hello 2" } ]
     });
     // After overriding the turn context, the environment context should be emitted again
-    // reflecting the new cwd, approval policy and sandbox settings.
-    let shell = default_user_shell().await;
-    let expected_env_text_2 = format!(
-        "<environment_context>\nCurrent working directory: {}\nApproval policy: never\nSandbox mode: workspace-write\nNetwork access: enabled\n{}</environment_context>",
-        new_cwd.path().to_string_lossy(),
-        match shell.name() {
-            Some(name) => format!("Shell: {name}\n"),
-            None => String::new(),
-        }
-    );
+    // reflecting the new approval policy and sandbox settings. Omit cwd because it did
+    // not change.
+    let expected_env_text_2 = r#"<environment_context>
+  <approval_policy>never</approval_policy>
+  <sandbox_mode>workspace-write</sandbox_mode>
+  <network_access>enabled</network_access>
+</environment_context>"#;
     let expected_env_msg_2 = serde_json::json!({
         "type": "message",
         "id": serde_json::Value::Null,