Skip to content
Merged
4 changes: 1 addition & 3 deletions codex-rs/core/src/arc_monitor_tests.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
use std::env;
use std::ffi::OsStr;
use std::path::PathBuf;
use std::sync::Arc;

use pretty_assertions::assert_eq;
Expand Down Expand Up @@ -74,8 +73,7 @@ async fn build_arc_monitor_request_includes_relevant_history_and_null_policies()
.record_into_history(
&[ContextualUserFragment::into(
crate::context::EnvironmentContext::new(
Some(PathBuf::from("/tmp")),
"zsh".to_string(),
Vec::new(),
/*current_date*/ None,
/*timezone*/ None,
/*network*/ None,
Expand Down
173 changes: 137 additions & 36 deletions codex-rs/core/src/context/environment_context.rs
Original file line number Diff line number Diff line change
@@ -1,21 +1,85 @@
use crate::session::turn_context::TurnContext;
use crate::shell::Shell;
use codex_protocol::protocol::TurnContextItem;
use codex_protocol::protocol::TurnContextNetworkItem;
use std::path::PathBuf;
use codex_utils_absolute_path::AbsolutePathBuf;

use super::ContextualUserFragment;

#[derive(Debug, Clone, PartialEq)]
pub(crate) struct EnvironmentContext {
pub(crate) cwd: Option<PathBuf>,
pub(crate) shell: String,
pub(crate) environments: EnvironmentContextEnvironments,
pub(crate) current_date: Option<String>,
pub(crate) timezone: Option<String>,
pub(crate) network: Option<NetworkContext>,
pub(crate) subagents: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct EnvironmentContextEnvironment {
pub(crate) id: String,
pub(crate) cwd: AbsolutePathBuf,
pub(crate) shell: String,
}

impl EnvironmentContextEnvironment {
fn legacy(cwd: AbsolutePathBuf, shell: String) -> Self {
Self {
id: String::new(),
cwd,
shell,
}
}

fn from_turn_environments(
environments: &[crate::session::turn_context::TurnEnvironment],
) -> Vec<Self> {
environments
.iter()
.map(|environment| Self {
id: environment.environment_id.clone(),
cwd: environment.cwd.clone(),
shell: environment.shell.clone(),
})
.collect()
}
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum EnvironmentContextEnvironments {
None,
Single(EnvironmentContextEnvironment),
Multiple(Vec<EnvironmentContextEnvironment>),
}

impl EnvironmentContextEnvironments {
fn from_vec(environments: Vec<EnvironmentContextEnvironment>) -> Self {
let mut environments = environments;
match environments.pop() {
None => Self::None,
Some(environment) if environments.is_empty() => Self::Single(environment),
Some(environment) => {
environments.push(environment);
Self::Multiple(environments)
}
}
}

fn equals_except_shell(&self, other: &Self) -> bool {
match (self, other) {
(Self::None, Self::None) => true,
(Self::Single(left), Self::Single(right)) => left.cwd == right.cwd,
(Self::Multiple(left), Self::Multiple(right)) => {
left.len() == right.len()
&& left
.iter()
.zip(right.iter())
.all(|(left, right)| left.id == right.id && left.cwd == right.cwd)
}
_ => false,
}
}
}

#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub(crate) struct NetworkContext {
allowed_domains: Vec<String>,
Expand All @@ -33,70 +97,86 @@ impl NetworkContext {

impl EnvironmentContext {
pub(crate) fn new(
cwd: Option<PathBuf>,
shell: String,
environments: Vec<EnvironmentContextEnvironment>,
current_date: Option<String>,
timezone: Option<String>,
network: Option<NetworkContext>,
subagents: Option<String>,
) -> Self {
Self {
cwd,
shell,
environments: EnvironmentContextEnvironments::from_vec(environments),
current_date,
timezone,
network,
subagents,
}
}

/// Compares two environment contexts, ignoring the shell. Useful when
/// comparing turn to turn, since the initial environment_context will
/// include the shell, and then it is not configurable from turn to turn.
pub(crate) fn equals_except_shell(&self, other: &EnvironmentContext) -> bool {
let EnvironmentContext {
cwd,
fn new_with_environments(
environments: EnvironmentContextEnvironments,
current_date: Option<String>,
timezone: Option<String>,
network: Option<NetworkContext>,
subagents: Option<String>,
) -> Self {
Self {
environments,
current_date,
timezone,
network,
subagents,
shell: _,
} = other;
self.cwd == *cwd
&& self.current_date == *current_date
&& self.timezone == *timezone
&& self.network == *network
&& self.subagents == *subagents
}
}

/// Compares two environment contexts, ignoring the shell. Useful when
/// comparing turn to turn, since the initial environment_context will
/// include the shell, and then it is not configurable from turn to turn.
pub(crate) fn equals_except_shell(&self, other: &EnvironmentContext) -> bool {
self.environments.equals_except_shell(&other.environments)
&& self.current_date == other.current_date
&& self.timezone == other.timezone
&& self.network == other.network
&& self.subagents == other.subagents
}

pub(crate) fn diff_from_turn_context_item(
before: &TurnContextItem,
after: &EnvironmentContext,
) -> Self {
let before_network = Self::network_from_turn_context_item(before);
let cwd = match &after.cwd {
Some(cwd) if before.cwd.as_path() != cwd.as_path() => Some(cwd.clone()),
_ => None,
let environments = match &after.environments {
EnvironmentContextEnvironments::Single(environment) => {
if before.cwd.as_path() != environment.cwd.as_path() {
EnvironmentContextEnvironments::Single(EnvironmentContextEnvironment::legacy(
environment.cwd.clone(),
environment.shell.clone(),
))
} else {
EnvironmentContextEnvironments::None
}
}
EnvironmentContextEnvironments::Multiple(environments) => {
EnvironmentContextEnvironments::Multiple(environments.clone())
}
EnvironmentContextEnvironments::None => EnvironmentContextEnvironments::None,
};
let network = if before_network != after.network {
after.network.clone()
} else {
before_network
};
EnvironmentContext::new(
cwd,
after.shell.clone(),
EnvironmentContext::new_with_environments(
environments,
after.current_date.clone(),
after.timezone.clone(),
network,
/*subagents*/ None,
)
}

pub(crate) fn from_turn_context(turn_context: &TurnContext, shell: &Shell) -> Self {
pub(crate) fn from_turn_context(turn_context: &TurnContext) -> Self {
Self::new(
Some(turn_context.cwd.to_path_buf()),
shell.name().to_string(),
EnvironmentContextEnvironment::from_turn_environments(&turn_context.environments),
turn_context.current_date.clone(),
turn_context.timezone.clone(),
Self::network_from_turn_context(turn_context),
Expand All @@ -108,9 +188,12 @@ impl EnvironmentContext {
turn_context_item: &TurnContextItem,
shell: String,
) -> Self {
let cwd = match AbsolutePathBuf::try_from(turn_context_item.cwd.clone()) {
Ok(cwd) => cwd,
Err(_) => AbsolutePathBuf::resolve_path_against_base(&turn_context_item.cwd, "/"),
};
Self::new(
Some(turn_context_item.cwd.clone()),
shell,
vec![EnvironmentContextEnvironment::legacy(cwd, shell)],
turn_context_item.current_date.clone(),
turn_context_item.timezone.clone(),
Self::network_from_turn_context_item(turn_context_item),
Expand Down Expand Up @@ -168,11 +251,29 @@ impl ContextualUserFragment for EnvironmentContext {

fn body(&self) -> String {
let mut lines = Vec::new();
if let Some(cwd) = &self.cwd {
lines.push(format!(" <cwd>{}</cwd>", cwd.to_string_lossy()));
match &self.environments {
EnvironmentContextEnvironments::Single(environment) => {
lines.push(format!(
" <cwd>{}</cwd>",
environment.cwd.to_string_lossy()
));
lines.push(format!(" <shell>{}</shell>", environment.shell));
}
EnvironmentContextEnvironments::Multiple(environments) => {
lines.push(" <environments>".to_string());
for environment in environments {
lines.push(format!(" <environment id=\"{}\">", environment.id));
lines.push(format!(
" <cwd>{}</cwd>",
environment.cwd.to_string_lossy()
));
lines.push(format!(" <shell>{}</shell>", environment.shell));
lines.push(" </environment>".to_string());
}
lines.push(" </environments>".to_string());
}
EnvironmentContextEnvironments::None => {}
}

lines.push(format!(" <shell>{}</shell>", self.shell));
if let Some(current_date) = &self.current_date {
lines.push(format!(" <current_date>{current_date}</current_date>"));
}
Expand Down
Loading
Loading