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

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;
}
}

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,
})
}

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()

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()),
])
);

View File

@@ -786,15 +786,6 @@ from pathlib import Path
env_file = Path(os.environ["CLAUDE_ENV_FILE"])
with env_file.open("a", encoding="utf-8") as handle:
handle.write("export CODEX_SESSION_START_TEST='from-session-start'\n")
"#;
let pre_tool_use_script_path = home.join("pre_tool_use_env_hook.py");
let pre_tool_use_script = r#"import os
from pathlib import Path
env_file = os.environ.get("CODEX_ENV_FILE") or os.environ.get("CLAUDE_ENV_FILE")
if env_file:
with Path(env_file).open("a", encoding="utf-8") as handle:
handle.write("export CODEX_SESSION_START_TEST='leaked-to-pre-tool-use'\n")
"#;
let hooks = serde_json::json!({
"hooks": {
@@ -804,22 +795,12 @@ if env_file:
"command": format!("python3 {}", session_start_script_path.display()),
"statusMessage": "exporting session environment",
}]
}],
"PreToolUse": [{
"matcher": "^Bash$",
"hooks": [{
"type": "command",
"command": format!("python3 {}", pre_tool_use_script_path.display()),
"statusMessage": "checking session environment isolation",
}]
}]
}
});
fs::write(&session_start_script_path, session_start_script)
.context("write session start env hook script")?;
fs::write(&pre_tool_use_script_path, pre_tool_use_script)
.context("write pre tool use env hook script")?;
fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?;
Ok(())
}
@@ -2776,10 +2757,6 @@ async fn assert_session_start_env_file_reaches_bash_surface(
output.contains("from-session-start||"),
"expected env-file paths to stay hidden from {surface:?} command: {output}"
);
assert!(
!output.contains("leaked-to-pre-tool-use"),
"PreToolUse should not receive env-file paths: {output}"
);
Ok(())
}
@@ -2794,62 +2771,6 @@ async fn session_start_env_file_exports_reach_exec_command() -> Result<()> {
assert_session_start_env_file_reaches_bash_surface(BashRewriteSurface::ExecCommand).await
}
#[tokio::test]
async fn shell_command_hides_env_file_paths_when_hooks_are_disabled() -> Result<()> {
skip_if_no_network!(Ok(()));
skip_if_windows!(Ok(()));
let server = start_mock_server().await;
let call_id = "codex-env-file-without-hooks";
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-1"),
ev_function_call(
call_id,
"shell_command",
&serde_json::to_string(&serde_json::json!({
"command": "printf '%s|%s' \"${CODEX_ENV_FILE:+configured}\" \"${CLAUDE_ENV_FILE:+configured}\"",
"login": false,
}))?,
),
ev_completed("resp-1"),
]),
sse(vec![
ev_response_created("resp-2"),
ev_assistant_message("msg-1", "session env file available"),
ev_completed("resp-2"),
]),
],
)
.await;
let mut builder = test_codex().with_config(|config| {
if let Err(error) = config.features.disable(Feature::CodexHooks) {
panic!("test config should allow feature update: {error}");
}
});
let test = builder.build(&server).await?;
test.submit_turn_with_permission_profile(
"check whether the session environment file is available",
PermissionProfile::Disabled,
)
.await?;
let requests = responses.requests();
assert_eq!(requests.len(), 2);
assert!(
!requests[1]
.function_call_output(call_id)
.to_string()
.contains("configured")
);
Ok(())
}
async fn assert_pre_tool_use_rewrites_bash_surface(surface: BashRewriteSurface) -> Result<()> {
skip_if_no_network!(Ok(()));

View File

@@ -121,7 +121,7 @@ fn build_command(shell: &CommandShell, handler: &ConfiguredHandler) -> Command {
command.env_remove(CODEX_ENV_FILE_ENV_VAR);
command.env_remove(CLAUDE_ENV_FILE_ENV_VAR);
if handler.event_name == HookEventName::SessionStart {
command.envs(&shell.env);
command.envs(&shell.session_start_env);
}
command
}
@@ -165,7 +165,7 @@ mod tests {
CommandShell {
program: "hook-shell".to_string(),
args: Vec::new(),
env: HashMap::from([
session_start_env: HashMap::from([
(
CODEX_ENV_FILE_ENV_VAR.to_string(),
"session-owned-env-file".to_string(),
@@ -207,18 +207,4 @@ mod tests {
assert_eq!(command_env(&command, CODEX_ENV_FILE_ENV_VAR), Some(None));
assert_eq!(command_env(&command, CLAUDE_ENV_FILE_ENV_VAR), Some(None));
}
#[test]
fn session_start_hook_receives_session_owned_env_file_paths() {
let command = build_command(&shell(), &handler(HookEventName::SessionStart));
assert_eq!(
command_env(&command, CODEX_ENV_FILE_ENV_VAR),
Some(Some(OsStr::new("session-owned-env-file")))
);
assert_eq!(
command_env(&command, CLAUDE_ENV_FILE_ENV_VAR),
Some(Some(OsStr::new("session-owned-env-file")))
);
}
}

View File

@@ -36,7 +36,7 @@ use std::collections::HashMap;
pub(crate) struct CommandShell {
pub program: String,
pub args: Vec<String>,
pub env: HashMap<String, String>,
pub session_start_env: HashMap<String, String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]

View File

@@ -203,7 +203,7 @@ with Path(r"{log_path}").open("a", encoding="utf-8") as handle:
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -219,7 +219,7 @@ with Path(r"{log_path}").open("a", encoding="utf-8") as handle:
plugin_hook_load_warnings: Vec::new(),
shell_program: None,
shell_args: Vec::new(),
command_env: HashMap::new(),
session_start_env: HashMap::new(),
});
assert!(listed.hooks[0].is_managed);
let cwd = cwd();
@@ -308,7 +308,7 @@ async fn requirements_managed_hooks_execute_windows_command_override() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -388,7 +388,7 @@ fn unknown_requirement_source_hooks_stay_managed() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -471,7 +471,7 @@ fn user_disablement_filters_non_managed_hooks_but_not_managed_hooks() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -535,7 +535,7 @@ fn user_disablement_does_not_filter_managed_layer_hooks() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -697,7 +697,7 @@ fn requirements_managed_hooks_load_when_managed_dir_is_missing() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -754,7 +754,7 @@ fn allow_managed_hooks_only_false_keeps_unmanaged_hooks() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -809,7 +809,7 @@ fn allow_managed_hooks_only_in_config_toml_does_not_enable_policy() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -880,7 +880,7 @@ fn allow_managed_hooks_only_skips_unmanaged_json_and_toml_hooks() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -920,7 +920,7 @@ fn allow_managed_hooks_only_skips_unmanaged_plugin_hooks() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -993,7 +993,7 @@ fn allow_managed_hooks_only_keeps_managed_requirement_and_config_layer_hooks() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -1104,7 +1104,7 @@ fn discovers_hooks_from_json_and_toml_in_the_same_layer() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -1198,7 +1198,7 @@ print(json.dumps({
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -1227,7 +1227,7 @@ print(json.dumps({
plugin_hook_load_warnings: Vec::new(),
shell_program: None,
shell_args: Vec::new(),
command_env: HashMap::new(),
session_start_env: HashMap::new(),
});
assert_eq!(
listed.hooks[0].plugin_id.as_deref(),
@@ -1314,7 +1314,7 @@ fn plugin_hook_sources_expand_plugin_placeholders() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);
@@ -1359,7 +1359,7 @@ fn plugin_hook_load_warnings_are_startup_warnings() {
CommandShell {
program: String::new(),
args: Vec::new(),
env: HashMap::new(),
session_start_env: HashMap::new(),
},
);

View File

@@ -182,8 +182,10 @@ pub(crate) async fn run(
};
let results = if event_name == HookEventName::SessionStart
&& (shell.env.contains_key(CODEX_ENV_FILE_ENV_VAR)
|| shell.env.contains_key(CLAUDE_ENV_FILE_ENV_VAR))
&& (shell.session_start_env.contains_key(CODEX_ENV_FILE_ENV_VAR)
|| shell
.session_start_env
.contains_key(CLAUDE_ENV_FILE_ENV_VAR))
{
// Serialize hooks that share the session env file to preserve configured order.
let mut results = Vec::with_capacity(matched.len());
@@ -544,7 +546,7 @@ mod tests {
let shell = CommandShell {
program: "sh".to_string(),
args: vec!["-c".to_string()],
env: HashMap::from([(
session_start_env: HashMap::from([(
CODEX_ENV_FILE_ENV_VAR.to_string(),
env_file.path().display().to_string(),
)]),

View File

@@ -37,7 +37,7 @@ pub struct HooksConfig {
pub plugin_hook_load_warnings: Vec<String>,
pub shell_program: Option<String>,
pub shell_args: Vec<String>,
pub command_env: HashMap<String, String>,
pub session_start_env: HashMap<String, String>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
@@ -75,7 +75,7 @@ impl Hooks {
CommandShell {
program: config.shell_program.unwrap_or_default(),
args: config.shell_args,
env: config.command_env,
session_start_env: config.session_start_env,
},
);
Self {