Address SessionStart env file review

This commit is contained in:
Abhinav Vedmala
2026-05-30 15:01:23 -07:00
parent 027bd73a4c
commit fac12e3895
10 changed files with 328 additions and 51 deletions
+13 -2
View File
@@ -3240,10 +3240,21 @@ impl Session {
pub(crate) fn apply_shell_env_exports(
&self,
env: &mut HashMap<String, String>,
explicit_env_overrides: &HashMap<String, String>,
policy: &ShellEnvironmentPolicy,
) {
if let Some(shell_env_file) = self.services.shell_env_file.as_ref() {
shell_env_file.apply_exports(env, explicit_env_overrides);
shell_env_file.apply_exports(env, policy);
}
}
pub(crate) fn extend_shell_env_snapshot_overrides(
&self,
overrides: &mut HashMap<String, String>,
env: &HashMap<String, String>,
policy: &ShellEnvironmentPolicy,
) {
if let Some(shell_env_file) = self.services.shell_env_file.as_ref() {
shell_env_file.extend_snapshot_overrides(overrides, env, policy);
}
}
+62 -13
View File
@@ -7,6 +7,8 @@ use anyhow::Result;
use anyhow::anyhow;
use anyhow::bail;
use codex_protocol::ThreadId;
use codex_protocol::config_types::EnvironmentVariablePattern;
use codex_protocol::config_types::ShellEnvironmentPolicy;
use codex_protocol::exec_output::bytes_to_string_smart;
use codex_protocol::shell_environment::CLAUDE_ENV_FILE_ENV_VAR;
use codex_protocol::shell_environment::CODEX_ENV_FILE_ENV_VAR;
@@ -109,12 +111,12 @@ impl ShellEnvFile {
pub(crate) fn apply_exports(
&self,
env: &mut HashMap<String, String>,
explicit_env_overrides: &HashMap<String, String>,
policy: &ShellEnvironmentPolicy,
) {
let thread_id = get_env_var(env, CODEX_THREAD_ID_ENV_VAR).cloned();
if let Ok(exports) = self.exports.lock() {
for (key, value) in exports.iter() {
if ignored_capture_key(key) {
if !should_apply_captured_export(key, policy) {
continue;
}
match value {
@@ -127,7 +129,7 @@ impl ShellEnvFile {
}
}
}
for (key, value) in explicit_env_overrides {
for (key, value) in &policy.r#set {
insert_env_var(env, key.clone(), value.clone());
}
if let Some(thread_id) = thread_id {
@@ -136,6 +138,26 @@ impl ShellEnvFile {
remove_env_var(env, CODEX_ENV_FILE_ENV_VAR);
remove_env_var(env, CLAUDE_ENV_FILE_ENV_VAR);
}
pub(crate) fn extend_snapshot_overrides(
&self,
overrides: &mut HashMap<String, String>,
env: &HashMap<String, String>,
policy: &ShellEnvironmentPolicy,
) {
if let Ok(exports) = self.exports.lock() {
for key in exports.keys() {
if !should_apply_captured_export(key, policy) {
continue;
}
if get_env_var(overrides, key).is_some() {
continue;
}
let value = get_env_var(env, key).cloned().unwrap_or_default();
overrides.insert(key.clone(), value);
}
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -218,22 +240,22 @@ impl ShellEnvCapture {
}
}
const POSIX_DUMP_ENV_SCRIPT: &str = r#"if [ -x /usr/bin/env ]; then
/usr/bin/env -0
else
env -0
fi"#;
const POSIX_DUMP_ENV_SCRIPT: &str = r#"awk 'BEGIN {
for (key in ENVIRON) {
printf "%s=%s%c", key, ENVIRON[key], 0
}
}'"#;
const POSIX_SOURCE_ENV_FILE_AND_DUMP_ENV_SCRIPT: &str = r#"if [ -n "${CODEX_ENV_FILE:-}" ] && [ -f "$CODEX_ENV_FILE" ]; then
if . "$CODEX_ENV_FILE" >/dev/null 2>&1; then
:
fi
fi
if [ -x /usr/bin/env ]; then
/usr/bin/env -0
else
env -0
fi"#;
awk 'BEGIN {
for (key in ENVIRON) {
printf "%s=%s%c", key, ENVIRON[key], 0
}
}'"#;
const POWERSHELL_DUMP_ENV_SCRIPT: &str = r#"$items = [ordered]@{}
Get-ChildItem Env: | Sort-Object Name | ForEach-Object {
@@ -336,6 +358,33 @@ fn ignored_capture_key(key: &str) -> bool {
.any(|ignored| env_key_eq(key, ignored))
}
fn should_apply_captured_export(key: &str, policy: &ShellEnvironmentPolicy) -> bool {
!ignored_capture_key(key) && captured_export_allowed_by_policy(key, policy)
}
fn captured_export_allowed_by_policy(key: &str, policy: &ShellEnvironmentPolicy) -> bool {
if !policy.ignore_default_excludes {
let default_excludes = [
EnvironmentVariablePattern::new_case_insensitive("*KEY*"),
EnvironmentVariablePattern::new_case_insensitive("*SECRET*"),
EnvironmentVariablePattern::new_case_insensitive("*TOKEN*"),
];
if matches_any_pattern(key, &default_excludes) {
return false;
}
}
if matches_any_pattern(key, &policy.exclude) {
return false;
}
policy.include_only.is_empty() || matches_any_pattern(key, &policy.include_only)
}
fn matches_any_pattern(key: &str, patterns: &[EnvironmentVariablePattern]) -> bool {
patterns.iter().any(|pattern| pattern.matches(key))
}
fn get_env_var<'a>(env: &'a HashMap<String, String>, key: &str) -> Option<&'a String> {
env.iter()
.find(|(candidate, _)| env_key_eq(candidate, key))
+107 -9
View File
@@ -1,5 +1,7 @@
use anyhow::Result;
use codex_protocol::ThreadId;
use codex_protocol::config_types::EnvironmentVariablePattern;
use codex_protocol::config_types::ShellEnvironmentPolicy;
use pretty_assertions::assert_eq;
use std::collections::HashMap;
use std::path::PathBuf;
@@ -61,10 +63,12 @@ export EXPLICIT_OVERRIDE='from-hook'
CLAUDE_ENV_FILE_ENV_VAR.to_string(),
env_file.path().display().to_string(),
);
let explicit_env_overrides =
HashMap::from([("EXPLICIT_OVERRIDE".to_string(), "from-policy".to_string())]);
let policy = ShellEnvironmentPolicy {
r#set: HashMap::from([("EXPLICIT_OVERRIDE".to_string(), "from-policy".to_string())]),
..Default::default()
};
env_file.apply_exports(&mut env, &explicit_env_overrides);
env_file.apply_exports(&mut env, &policy);
assert_eq!(
env,
@@ -90,6 +94,96 @@ export EXPLICIT_OVERRIDE='from-hook'
Ok(())
}
#[cfg(not(windows))]
#[tokio::test]
async fn shell_env_file_filters_captured_exports_before_applying() -> Result<()> {
let env_file = ShellEnvFile::new(ThreadId::new(), ShellEnvCapture::Posix)?;
let base_env = HashMap::from([("PATH".to_string(), "/usr/bin".to_string())]);
std::fs::write(
env_file.path(),
"\
export PATH=\"/plugin/bin:$PATH\"
export ALLOWED_VALUE='allowed'
export OPENAI_API_KEY='secret'
export BLOCKED_VALUE='blocked'
export NOT_INCLUDED='not-included'
export EXPLICIT_OVERRIDE='from-hook'
",
)?;
let cwd = std::env::current_dir()?;
env_file
.capture_exports(&test_shell(), cwd.as_path(), &base_env)
.await?;
let mut env = base_env;
let policy = ShellEnvironmentPolicy {
ignore_default_excludes: false,
exclude: vec![EnvironmentVariablePattern::new_case_insensitive(
"BLOCKED_*",
)],
include_only: vec![
EnvironmentVariablePattern::new_case_insensitive("PATH"),
EnvironmentVariablePattern::new_case_insensitive("ALLOWED_*"),
EnvironmentVariablePattern::new_case_insensitive("EXPLICIT_*"),
],
r#set: HashMap::from([("EXPLICIT_OVERRIDE".to_string(), "from-policy".to_string())]),
..Default::default()
};
env_file.apply_exports(&mut env, &policy);
assert_eq!(
env,
HashMap::from([
("PATH".to_string(), "/plugin/bin:/usr/bin".to_string()),
("ALLOWED_VALUE".to_string(), "allowed".to_string()),
("EXPLICIT_OVERRIDE".to_string(), "from-policy".to_string()),
])
);
Ok(())
}
#[cfg(not(windows))]
#[tokio::test]
async fn shell_env_file_snapshot_overrides_include_captured_export_keys() -> Result<()> {
let env_file = ShellEnvFile::new(ThreadId::new(), ShellEnvCapture::Posix)?;
let base_env = HashMap::from([
("PATH".to_string(), "/usr/bin".to_string()),
("REMOVED_BY_HOOK".to_string(), "remove-me".to_string()),
]);
std::fs::write(
env_file.path(),
"\
export PATH=\"/plugin/bin:$PATH\"
export HOOK_ONLY='from-hook'
unset REMOVED_BY_HOOK
",
)?;
let cwd = std::env::current_dir()?;
env_file
.capture_exports(&test_shell(), cwd.as_path(), &base_env)
.await?;
let mut env = base_env;
let policy = ShellEnvironmentPolicy::default();
env_file.apply_exports(&mut env, &policy);
let mut overrides = HashMap::new();
env_file.extend_snapshot_overrides(&mut overrides, &env, &policy);
assert_eq!(
overrides,
HashMap::from([
("PATH".to_string(), "/plugin/bin:/usr/bin".to_string()),
("HOOK_ONLY".to_string(), "from-hook".to_string()),
("REMOVED_BY_HOOK".to_string(), String::new()),
])
);
Ok(())
}
#[cfg(not(windows))]
fn test_shell() -> Shell {
Shell {
@@ -191,10 +285,12 @@ Remove-Item Env:REMOVED_BY_HOOK -ErrorAction SilentlyContinue
CLAUDE_ENV_FILE_ENV_VAR.to_string(),
env_file.path().display().to_string(),
);
let explicit_env_overrides =
HashMap::from([("EXPLICIT_OVERRIDE".to_string(), "from-policy".to_string())]);
let policy = ShellEnvironmentPolicy {
r#set: HashMap::from([("EXPLICIT_OVERRIDE".to_string(), "from-policy".to_string())]),
..Default::default()
};
env_file.apply_exports(&mut env, &explicit_env_overrides);
env_file.apply_exports(&mut env, &policy);
assert_eq!(
env,
@@ -254,10 +350,12 @@ set REMOVED_BY_HOOK=\r
CLAUDE_ENV_FILE_ENV_VAR.to_string(),
env_file.path().display().to_string(),
);
let explicit_env_overrides =
HashMap::from([("EXPLICIT_OVERRIDE".to_string(), "from-policy".to_string())]);
let policy = ShellEnvironmentPolicy {
r#set: HashMap::from([("EXPLICIT_OVERRIDE".to_string(), "from-policy".to_string())]),
..Default::default()
};
env_file.apply_exports(&mut env, &explicit_env_overrides);
env_file.apply_exports(&mut env, &policy);
assert_eq!(
env,
+8 -5
View File
@@ -130,10 +130,7 @@ pub(crate) async fn execute_user_shell_command(
&turn_context.shell_environment_policy,
Some(session.conversation_id),
);
session.apply_shell_env_exports(
&mut exec_env_map,
&turn_context.shell_environment_policy.r#set,
);
session.apply_shell_env_exports(&mut exec_env_map, &turn_context.shell_environment_policy);
if exec_env_map.contains_key(PROXY_ACTIVE_ENV_KEY) {
for key in PROXY_ENV_KEYS {
exec_env_map.remove(*key);
@@ -148,12 +145,18 @@ pub(crate) async fn execute_user_shell_command(
exec_env_map.remove(PROXY_GIT_SSH_COMMAND_ENV_KEY);
}
}
let mut snapshot_overrides = turn_context.shell_environment_policy.r#set.clone();
session.extend_shell_env_snapshot_overrides(
&mut snapshot_overrides,
&exec_env_map,
&turn_context.shell_environment_policy,
);
let exec_command = maybe_wrap_shell_lc_with_snapshot(
&display_command,
session_shell.as_ref(),
#[allow(deprecated)]
&turn_context.cwd,
&turn_context.shell_environment_policy.r#set,
&snapshot_overrides,
&exec_env_map,
);
@@ -96,7 +96,7 @@ impl ShellCommandHandler {
let cwd = turn_context.resolve_path(params.workdir.clone());
let mut env = create_env(&turn_context.shell_environment_policy, Some(thread_id));
session.apply_shell_env_exports(&mut env, &turn_context.shell_environment_policy.r#set);
session.apply_shell_env_exports(&mut env, &turn_context.shell_environment_policy);
Ok(ExecParams {
command,
@@ -248,6 +248,12 @@ impl ToolRuntime<ShellRequest, ExecToolCallOutput> for ShellRuntime {
}
(env, explicit_env_overrides)
};
let mut explicit_env_overrides = explicit_env_overrides;
ctx.session.extend_shell_env_snapshot_overrides(
&mut explicit_env_overrides,
&env,
&ctx.turn.shell_environment_policy,
);
let command = maybe_wrap_shell_lc_with_snapshot(
&req.command,
session_shell.as_ref(),
@@ -277,6 +277,14 @@ impl<'a> ToolRuntime<UnifiedExecRequest, UnifiedExecProcess> for UnifiedExecRunt
explicit_env_overrides
};
let environment_is_remote = req.environment.is_remote();
let mut explicit_env_overrides = explicit_env_overrides;
if !environment_is_remote {
ctx.session.extend_shell_env_snapshot_overrides(
&mut explicit_env_overrides,
&env,
&ctx.turn.shell_environment_policy,
);
}
let command = if environment_is_remote {
base_command.to_vec()
} else {
@@ -1007,7 +1007,7 @@ impl UnifiedExecProcessManager {
if !request.environment.is_remote() {
context
.session
.apply_shell_env_exports(&mut env, &context.turn.shell_environment_policy.r#set);
.apply_shell_env_exports(&mut env, &context.turn.shell_environment_policy);
}
let env = apply_unified_exec_env(env);
let exec_server_env_config = ExecServerEnvConfig {
+76 -2
View File
@@ -631,9 +631,12 @@ printf 'export FIRST=1\n' >> "$CODEX_ENV_FILE""#
]),
};
let mut hook = handler();
hook.command = r#"[ "$CODEX_ENV_FILE" = "$CLAUDE_ENV_FILE" ] || exit 9
hook.command =
r#"mode=$(stat -f %Lp "$CODEX_ENV_FILE" 2>/dev/null || stat -c %a "$CODEX_ENV_FILE")
[ "$mode" = 600 ] || exit 8
[ "$CODEX_ENV_FILE" = "$CLAUDE_ENV_FILE" ] || exit 9
printf 'export ALIAS_OK=1\n' >> "$CLAUDE_ENV_FILE""#
.to_string();
.to_string();
let cwd = tempdir().expect("create cwd");
let outcome = run(
@@ -667,6 +670,77 @@ printf 'export ALIAS_OK=1\n' >> "$CLAUDE_ENV_FILE""#
);
}
#[cfg(not(windows))]
#[tokio::test]
async fn session_start_hooks_merge_only_successful_env_files_with_separator() {
use std::collections::HashMap;
use std::fs;
use codex_protocol::ThreadId;
use codex_protocol::shell_environment::CODEX_ENV_FILE_ENV_VAR;
use tempfile::NamedTempFile;
use tempfile::tempdir;
use super::SessionStartRequest;
use super::SessionStartSource;
use super::StartHookTarget;
use super::run;
use crate::engine::CommandShell;
let env_file = NamedTempFile::new().expect("create env file");
let shell = CommandShell {
program: "sh".to_string(),
args: vec!["-c".to_string()],
session_start_env: HashMap::from([(
CODEX_ENV_FILE_ENV_VAR.to_string(),
env_file.path().display().to_string(),
)]),
};
let mut first = handler();
first.command = r#"printf 'export FIRST=1' >> "$CODEX_ENV_FILE""#.to_string();
let mut failed = handler();
failed.command = r#"printf 'export FAILED=1\n' >> "$CODEX_ENV_FILE"; exit 3"#.to_string();
failed.display_order = 1;
let mut third = handler();
third.command = r#"printf 'export THIRD=1\n' >> "$CODEX_ENV_FILE""#.to_string();
third.display_order = 2;
let cwd = tempdir().expect("create cwd");
let outcome = run(
&[first, failed, third],
&shell,
SessionStartRequest {
session_id: ThreadId::new(),
cwd: cwd.path().to_path_buf().abs(),
transcript_path: None,
model: "gpt-test".to_string(),
permission_mode: "default".to_string(),
target: StartHookTarget::SessionStart {
source: SessionStartSource::Startup,
},
},
/*turn_id*/ None,
)
.await;
assert_eq!(
outcome
.hook_events
.iter()
.map(|event| event.run.status)
.collect::<Vec<_>>(),
vec![
HookRunStatus::Completed,
HookRunStatus::Failed,
HookRunStatus::Completed
]
);
assert_eq!(
fs::read_to_string(env_file.path()).expect("read env file"),
"export FIRST=1\nexport THIRD=1\n"
);
}
fn handler() -> ConfiguredHandler {
handler_for(HookEventName::SessionStart)
}
@@ -1,4 +1,6 @@
use std::io::ErrorKind;
#[cfg(unix)]
use std::os::unix::fs::OpenOptionsExt;
use std::path::Path;
use std::path::PathBuf;
@@ -60,17 +62,16 @@ pub(super) async fn execute_handlers<T>(
Uuid::new_v4()
));
let scratch_result = async {
fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(&scratch_path)
.await
.with_context(|| {
format!(
"failed to create SessionStart scratch env file {}",
scratch_path.display()
)
})?;
let mut options = std::fs::OpenOptions::new();
options.write(true).create_new(true);
#[cfg(unix)]
options.mode(0o600);
options.open(&scratch_path).with_context(|| {
format!(
"failed to create SessionStart scratch env file {}",
scratch_path.display()
)
})?;
let mut scratch_shell = shell.clone();
let scratch_path_env = scratch_path.to_string_lossy().to_string();
@@ -132,13 +133,30 @@ pub(super) async fn execute_handlers<T>(
}
completed.sort_by_key(|execution| execution.configured_order);
// Merge each handler's isolated file in configured order. The core session
// layer will source this single canonical file after hooks finish.
let mut env_file_ends_with_newline = match fs::read(&env_file_path).await {
Ok(contents) => contents.last().is_none_or(|byte| *byte == b'\n'),
Err(error) if error.kind() == ErrorKind::NotFound => true,
Err(error) => {
tracing::warn!(
"failed to read SessionStart env file {} before merge: {error}",
env_file_path.display()
);
true
}
};
// Merge each successful handler's isolated file in configured order. The
// core session layer will source this single canonical file after hooks
// finish.
for execution in &mut completed {
let Some(scratch_path) = execution.scratch_path.take() else {
continue;
};
let result: Result<()> = async {
let result: Result<bool> = async {
if execution.run_result.error.is_some() || execution.run_result.exit_code != Some(0) {
return Ok(env_file_ends_with_newline);
}
let contents = fs::read(&scratch_path).await.with_context(|| {
format!(
"failed to read SessionStart scratch env file {}",
@@ -146,7 +164,7 @@ pub(super) async fn execute_handlers<T>(
)
})?;
if contents.is_empty() {
return Ok(());
return Ok(env_file_ends_with_newline);
}
let mut file = fs::OpenOptions::new()
@@ -155,13 +173,18 @@ pub(super) async fn execute_handlers<T>(
.open(&env_file_path)
.await
.with_context(|| format!("failed to open {}", env_file_path.display()))?;
if !env_file_ends_with_newline && contents.first() != Some(&b'\n') {
file.write_all(b"\n")
.await
.with_context(|| format!("failed to append to {}", env_file_path.display()))?;
}
file.write_all(&contents)
.await
.with_context(|| format!("failed to append to {}", env_file_path.display()))?;
file.flush()
.await
.with_context(|| format!("failed to flush {}", env_file_path.display()))?;
Ok(())
Ok(contents.last() == Some(&b'\n'))
}
.await;
@@ -174,8 +197,13 @@ pub(super) async fn execute_handlers<T>(
scratch_path.display()
);
}
if let Err(error) = result {
execution.run_result = error_result(error);
match result {
Ok(ends_with_newline) => {
env_file_ends_with_newline = ends_with_newline;
}
Err(error) => {
execution.run_result = error_result(error);
}
}
}