[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
This commit is contained in:
Dylan
2025-08-20 23:45:16 -07:00
committed by GitHub
parent 74683bab91
commit d2b2a6d13a
3 changed files with 85 additions and 57 deletions

View File

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

View File

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

View File

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