Simplify session env file handling

This commit is contained in:
Abhinav Vedmala
2026-05-28 23:09:12 -07:00
parent 09c8e31c90
commit 59d6d03cbd
10 changed files with 82 additions and 220 deletions
+10 -12
View File
@@ -125,7 +125,6 @@ pub(crate) async fn run_pending_session_start_hooks(
source: session_start_source,
},
};
let captures_shell_env = matches!(target, StartHookTarget::SessionStart { .. });
let request = codex_hooks::SessionStartRequest {
session_id: sess.session_id().into(),
#[allow(deprecated)]
@@ -137,10 +136,15 @@ pub(crate) async fn run_pending_session_start_hooks(
};
let hooks = sess.hooks();
let preview_runs = hooks.preview_session_start(&request);
emit_hook_started_events(sess, turn_context, preview_runs).await;
let outcome = hooks
.run_session_start(request, Some(turn_context.sub_id.clone()))
.await;
let captures_shell_env = matches!(&request.target, StartHookTarget::SessionStart { .. })
&& !preview_runs.is_empty();
let outcome = run_context_injecting_hook(
sess,
turn_context,
preview_runs,
hooks.run_session_start(request, Some(turn_context.sub_id.clone())),
)
.await;
if captures_shell_env {
let base_env = create_env(
&turn_context.shell_environment_policy,
@@ -155,13 +159,7 @@ pub(crate) async fn run_pending_session_start_hooks(
tracing::warn!("failed to capture SessionStart shell environment: {error:#}");
}
}
let outcome: ContextInjectingHookOutcome = outcome.into();
emit_hook_completed_events(sess, turn_context, outcome.hook_events).await;
if outcome
.outcome
.record_additional_contexts(sess, turn_context)
.await
{
if outcome.record_additional_contexts(sess, turn_context).await {
return true;
}
}
+3 -3
View File
@@ -3326,9 +3326,9 @@ async fn build_hooks_for_config(
let plugin_outcome = plugins_manager.plugins_for_config(&plugins_input).await;
let plugin_hook_sources = plugin_outcome.effective_plugin_hook_sources();
let plugin_hook_load_warnings = plugin_outcome.effective_plugin_hook_warnings();
let mut command_env = HashMap::new();
let mut session_start_env = HashMap::new();
if let Some(shell_env_file) = shell_env_file {
shell_env_file.insert_path_into_env(&mut command_env);
shell_env_file.insert_path_into_env(&mut session_start_env);
}
Hooks::new(HooksConfig {
legacy_notify_argv: config.notify.clone(),
@@ -3339,7 +3339,7 @@ async fn build_hooks_for_config(
plugin_hook_load_warnings,
shell_program: Some(hook_shell_program),
shell_args: hook_shell_argv,
command_env,
session_start_env,
})
}
+33 -50
View File
@@ -72,7 +72,7 @@ impl ShellEnvFile {
cwd: &Path,
base_env: &HashMap<String, String>,
) -> Result<()> {
if !self.capture.supports_shell_type(&shell.shell_type) {
if ShellEnvCapture::for_shell_type(&shell.shell_type) != Some(self.capture) {
return Ok(());
}
@@ -81,13 +81,14 @@ impl ShellEnvFile {
remove_env_var(&mut capture_env, CLAUDE_ENV_FILE_ENV_VAR);
self.insert_path_into_env(&mut capture_env);
let (baseline_script, source_script) = self.capture.scripts();
let baseline = self
.capture
.capture_env_from_shell(shell, cwd, ShellEnvCaptureMode::Baseline, &capture_env)
.capture_env_from_shell(shell, cwd, baseline_script, &capture_env)
.await?;
let captured = self
.capture
.capture_env_from_shell(shell, cwd, ShellEnvCaptureMode::SourceEnvFile, &capture_env)
.capture_env_from_shell(shell, cwd, source_script, &capture_env)
.await?;
let exports = diff_env(&baseline, &captured);
*self
@@ -110,21 +111,18 @@ impl ShellEnvFile {
explicit_env_overrides: &HashMap<String, String>,
) {
let thread_id = get_env_var(env, CODEX_THREAD_ID_ENV_VAR).cloned();
let exports = self
.exports
.lock()
.map(|exports| exports.clone())
.unwrap_or_default();
for (key, value) in exports {
if ignored_capture_key(&key) {
continue;
}
match value {
Some(value) => {
insert_env_var(env, key, value);
if let Ok(exports) = self.exports.lock() {
for (key, value) in exports.iter() {
if ignored_capture_key(key) {
continue;
}
None => {
remove_env_var(env, &key);
match value {
Some(value) => {
insert_env_var(env, key.clone(), value.clone());
}
None => {
remove_env_var(env, key);
}
}
}
}
@@ -145,12 +143,6 @@ enum ShellEnvCapture {
PowerShell,
}
#[derive(Clone, Copy, Debug)]
enum ShellEnvCaptureMode {
Baseline,
SourceEnvFile,
}
impl ShellEnvCapture {
fn for_shell_type(shell_type: &ShellType) -> Option<Self> {
match shell_type {
@@ -160,10 +152,6 @@ impl ShellEnvCapture {
}
}
fn supports_shell_type(self, shell_type: &ShellType) -> bool {
Self::for_shell_type(shell_type) == Some(self)
}
fn file_suffix(self) -> &'static str {
match self {
Self::Posix => ".sh",
@@ -171,14 +159,26 @@ impl ShellEnvCapture {
}
}
fn scripts(self) -> (&'static str, &'static str) {
match self {
Self::Posix => (
POSIX_DUMP_ENV_SCRIPT,
POSIX_SOURCE_ENV_FILE_AND_DUMP_ENV_SCRIPT,
),
Self::PowerShell => (
POWERSHELL_DUMP_ENV_SCRIPT,
POWERSHELL_SOURCE_ENV_FILE_AND_DUMP_ENV_SCRIPT,
),
}
}
async fn capture_env_from_shell(
self,
shell: &Shell,
cwd: &Path,
mode: ShellEnvCaptureMode,
script: &str,
env: &HashMap<String, String>,
) -> Result<HashMap<String, String>> {
let script = self.capture_script(mode);
let output = Command::new(&shell.shell_path)
.current_dir(cwd)
.args(self.capture_args(script))
@@ -197,19 +197,6 @@ impl ShellEnvCapture {
self.parse_env_output(&output.stdout)
}
fn capture_script(self, mode: ShellEnvCaptureMode) -> &'static str {
match (self, mode) {
(Self::Posix, ShellEnvCaptureMode::Baseline) => POSIX_DUMP_ENV_SCRIPT,
(Self::Posix, ShellEnvCaptureMode::SourceEnvFile) => {
POSIX_SOURCE_ENV_FILE_AND_DUMP_ENV_SCRIPT
}
(Self::PowerShell, ShellEnvCaptureMode::Baseline) => POWERSHELL_DUMP_ENV_SCRIPT,
(Self::PowerShell, ShellEnvCaptureMode::SourceEnvFile) => {
POWERSHELL_SOURCE_ENV_FILE_AND_DUMP_ENV_SCRIPT
}
}
}
fn capture_args(self, script: &str) -> Vec<&str> {
match self {
Self::Posix => vec!["-c", script],
@@ -270,11 +257,11 @@ fn parse_posix_env_output(output: &[u8]) -> Result<HashMap<String, String>> {
let Some(separator) = entry.iter().position(|byte| *byte == b'=') else {
continue;
};
let key = String::from_utf8(entry[..separator].to_vec())
let key = std::str::from_utf8(&entry[..separator])
.context("captured shell environment key was not UTF-8")?;
let value = String::from_utf8(entry[separator + 1..].to_vec())
let value = std::str::from_utf8(&entry[separator + 1..])
.context("captured shell environment value was not UTF-8")?;
env.insert(key, value);
env.insert(key.to_string(), value.to_string());
}
Ok(env)
}
@@ -308,7 +295,7 @@ fn diff_env(
if ignored_capture_key(key) {
continue;
}
if !contains_env_var(captured, key) {
if get_env_var(captured, key).is_none() {
exports.insert(key.clone(), None);
}
}
@@ -335,10 +322,6 @@ fn get_env_var<'a>(env: &'a HashMap<String, String>, key: &str) -> Option<&'a St
.map(|(_, value)| value)
}
fn contains_env_var(env: &HashMap<String, String>, key: &str) -> bool {
get_env_var(env, key).is_some()
}
fn insert_env_var(env: &mut HashMap<String, String>, key: String, value: String) {
if let Some(existing) = env
.keys()
+9 -37
View File
@@ -24,7 +24,7 @@ fn shell_env_file_is_removed_when_session_owner_drops() -> Result<()> {
#[cfg(not(windows))]
#[tokio::test]
async fn shell_env_file_applies_exports_without_exposing_writable_path() -> Result<()> {
async fn shell_env_file_captures_exports_without_exposing_writable_path() -> Result<()> {
let env_file = ShellEnvFile::new(ThreadId::new(), ShellEnvCapture::Posix)?;
let base_env = HashMap::from([
("PATH".to_string(), "/usr/bin".to_string()),
@@ -36,8 +36,11 @@ async fn shell_env_file_applies_exports_without_exposing_writable_path() -> Resu
std::fs::write(
env_file.path(),
"\
echo hidden
export CODEX_SESSION_START_TEST='from-session-start'
export PATH=\"/plugin/bin:$PATH\"
export COMMAND_SUBSTITUTION=$(printf unsafe)
export FUNCTION_DEF='() { echo unsafe; }'
export CODEX_ENV_FILE='/tmp/poison'
export CLAUDE_ENV_FILE='/tmp/poison'
export CODEX_THREAD_ID='poisoned-thread'
@@ -71,47 +74,16 @@ export EXPLICIT_OVERRIDE='from-hook'
"CODEX_SESSION_START_TEST".to_string(),
"from-session-start".to_string(),
),
(
CODEX_THREAD_ID_ENV_VAR.to_string(),
"real-thread".to_string(),
),
("EXPLICIT_OVERRIDE".to_string(), "from-policy".to_string()),
])
);
Ok(())
}
#[cfg(not(windows))]
#[tokio::test]
async fn shell_env_file_sources_shell_code_once() -> Result<()> {
let env_file = ShellEnvFile::new(ThreadId::new(), ShellEnvCapture::Posix)?;
std::fs::write(
env_file.path(),
"\
echo hidden
export SAFE=value
export COMMAND_SUBSTITUTION=$(printf unsafe)
export FUNCTION_DEF='() { echo unsafe; }'
",
)?;
let cwd = std::env::current_dir()?;
env_file
.capture_exports(&test_shell(), cwd.as_path(), &HashMap::new())
.await?;
let mut env = HashMap::new();
env_file.apply_exports(&mut env, &HashMap::new());
assert_eq!(
env,
HashMap::from([
("SAFE".to_string(), "value".to_string()),
("COMMAND_SUBSTITUTION".to_string(), "unsafe".to_string()),
(
"FUNCTION_DEF".to_string(),
"() { echo unsafe; }".to_string(),
),
(
CODEX_THREAD_ID_ENV_VAR.to_string(),
"real-thread".to_string(),
),
("EXPLICIT_OVERRIDE".to_string(), "from-policy".to_string()),
])
);