From d2b2a6d13ae4169a5806e8dc98c1333f188d894a Mon Sep 17 00:00:00 2001 From: Dylan Date: Wed, 20 Aug 2025 23:45:16 -0700 Subject: [PATCH] [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 --- codex-rs/core/src/codex.rs | 17 +++-- codex-rs/core/src/environment_context.rs | 94 +++++++++++++++--------- codex-rs/core/tests/prompt_caching.rs | 31 ++++---- 3 files changed, 85 insertions(+), 57 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index dfa10398..a692e6fa 100644 --- a/codex-rs/core/src/codex.rs +++ b/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; } diff --git a/codex-rs/core/src/environment_context.rs b/codex-rs/core/src/environment_context.rs index a5a9b855..c89d7ca6 100644 --- a/codex-rs/core/src/environment_context.rs +++ b/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 = "\n"; +pub(crate) const ENVIRONMENT_CONTEXT_START: &str = ""; pub(crate) const ENVIRONMENT_CONTEXT_END: &str = ""; #[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, + pub approval_policy: Option, + pub sandbox_mode: Option, + pub network_access: Option, + pub shell: Option, } impl EnvironmentContext { pub fn new( - cwd: PathBuf, - approval_policy: AskForApproval, - sandbox_policy: SandboxPolicy, - shell: Shell, + cwd: Option, + approval_policy: Option, + sandbox_policy: Option, + shell: Option, ) -> 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 + /// + /// ... + /// ... + /// ... + /// ... + /// ... + /// + /// ``` + 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.to_string_lossy())); } - Ok(()) + if let Some(approval_policy) = self.approval_policy { + lines.push(format!( + " {}", + approval_policy + )); + } + if let Some(sandbox_mode) = self.sandbox_mode { + lines.push(format!(" {}", sandbox_mode)); + } + if let Some(network_access) = self.network_access { + lines.push(format!( + " {}", + network_access + )); + } + if let Some(shell) = self.shell + && let Some(shell_name) = shell.name() + { + lines.push(format!(" {}", shell_name)); + } + lines.push(ENVIRONMENT_CONTEXT_END.to_string()); + lines.join("\n") } } @@ -86,7 +114,7 @@ impl From 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(), }], } } diff --git a/codex-rs/core/tests/prompt_caching.rs b/codex-rs/core/tests/prompt_caching.rs index 75b7691b..bf49262a 100644 --- a/codex-rs/core/tests/prompt_caching.rs +++ b/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!( - "\nCurrent working directory: {}\nApproval policy: on-request\nSandbox mode: read-only\nNetwork access: restricted\n{}", + r#" + {} + on-request + read-only + restricted +{}"#, cwd.path().to_string_lossy(), match shell.name() { - Some(name) => format!("Shell: {name}\n"), + Some(name) => format!(" {}\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::().unwrap(); let body2 = requests[1].body_json::().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!( - "\nCurrent working directory: {}\nApproval policy: never\nSandbox mode: workspace-write\nNetwork access: enabled\n{}", - 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#" + never + workspace-write + enabled +"#; let expected_env_msg_2 = serde_json::json!({ "type": "message", "id": serde_json::Value::Null,