Harden SessionStart env file lifecycle

This commit is contained in:
Abhinav Vedmala
2026-05-30 19:36:09 -07:00
parent fb5acdbaae
commit 83153b139f
9 changed files with 213 additions and 29 deletions

View File

@@ -136,8 +136,14 @@ pub(crate) async fn run_pending_session_start_hooks(
};
let hooks = sess.hooks();
let preview_runs = hooks.preview_session_start(&request);
let captures_shell_env = matches!(&request.target, StartHookTarget::SessionStart { .. })
&& !preview_runs.is_empty();
let is_session_start = matches!(&request.target, StartHookTarget::SessionStart { .. });
let captures_shell_env = is_session_start && !preview_runs.is_empty();
if is_session_start
&& !captures_shell_env
&& let Err(error) = sess.reset_shell_env_exports()
{
tracing::warn!("failed to reset SessionStart shell environment: {error:#}");
}
let outcome = run_context_injecting_hook(
sess,
turn_context,
@@ -157,6 +163,9 @@ pub(crate) async fn run_pending_session_start_hooks(
.unwrap_or_else(|| turn_context.config.cwd.as_path());
if let Err(error) = sess.capture_shell_env_exports(cwd, &base_env).await {
tracing::warn!("failed to capture SessionStart shell environment: {error:#}");
if let Err(error) = sess.reset_shell_env_exports() {
tracing::warn!("failed to reset SessionStart shell environment: {error:#}");
}
}
}
if outcome.record_additional_contexts(sess, turn_context).await {

View File

@@ -3237,6 +3237,13 @@ impl Session {
Ok(())
}
pub(crate) fn reset_shell_env_exports(&self) -> anyhow::Result<()> {
if let Some(shell_env_file) = self.services.shell_env_file.as_ref() {
shell_env_file.reset_exports()?;
}
Ok(())
}
pub(crate) fn apply_shell_env_exports(
&self,
env: &mut HashMap<String, String>,

View File

@@ -1,6 +1,7 @@
use std::collections::HashMap;
use std::path::Path;
use std::sync::Mutex;
use std::time::Duration;
use anyhow::Context;
use anyhow::Result;
@@ -14,11 +15,15 @@ use codex_protocol::shell_environment::CLAUDE_ENV_FILE_ENV_VAR;
use codex_protocol::shell_environment::CODEX_ENV_FILE_ENV_VAR;
use codex_protocol::shell_environment::CODEX_THREAD_ID_ENV_VAR;
use tempfile::TempPath;
use tokio::fs;
use tokio::process::Command;
use tokio::time::timeout;
use crate::shell::Shell;
use crate::shell::ShellType;
const SHELL_ENV_CAPTURE_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 10);
/// Session-owned script that hooks can populate with exported shell state.
///
/// Only lifecycle hooks receive the writable file path. After SessionStart
@@ -57,6 +62,14 @@ impl ShellEnvFile {
self.path.as_ref()
}
pub(crate) fn reset_exports(&self) -> Result<()> {
self.exports
.lock()
.map_err(|_| anyhow!("shell env exports lock poisoned"))?
.clear();
Ok(())
}
/// Sources the hook-writable env file once and stores the resulting
/// environment diff.
///
@@ -70,6 +83,22 @@ impl ShellEnvFile {
shell: &Shell,
cwd: &Path,
base_env: &HashMap<String, String>,
) -> Result<()> {
let result = self.capture_exports_inner(shell, cwd, base_env).await;
if let Err(error) = fs::write(self.path(), []).await {
tracing::warn!(
"failed to clear SessionStart shell env file {}: {error}",
self.path().display()
);
}
result
}
async fn capture_exports_inner(
&self,
shell: &Shell,
cwd: &Path,
base_env: &HashMap<String, String>,
) -> Result<()> {
let mut capture_env = base_env.clone();
remove_env_var(&mut capture_env, CODEX_ENV_FILE_ENV_VAR);
@@ -128,7 +157,9 @@ impl ShellEnvFile {
}
}
for (key, value) in &policy.r#set {
insert_env_var(env, key.clone(), value.clone());
if shell_environment::env_var_allowed_by_include_only(key, policy) {
insert_env_var(env, key.clone(), value.clone());
}
}
if let Some(thread_id) = thread_id {
insert_env_var(env, CODEX_THREAD_ID_ENV_VAR.to_string(), thread_id);
@@ -195,13 +226,21 @@ impl ShellEnvCapture {
script: &str,
env: &HashMap<String, String>,
) -> Result<HashMap<String, String>> {
let output = Command::new(&shell.shell_path)
let mut command = Command::new(&shell.shell_path);
command
.current_dir(cwd)
.args(self.capture_args(script))
.env_clear()
.envs(env)
.output()
.kill_on_drop(true);
let output = timeout(SHELL_ENV_CAPTURE_TIMEOUT, command.output())
.await
.with_context(|| {
format!(
"timed out capturing shell environment with {}",
shell.shell_path.display()
)
})?
.with_context(|| format!("failed to run {}", shell.shell_path.display()))?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
@@ -230,22 +269,14 @@ impl ShellEnvCapture {
}
}
const POSIX_DUMP_ENV_SCRIPT: &str = r#"awk 'BEGIN {
for (key in ENVIRON) {
printf "%s=%s%c", key, ENVIRON[key], 0
}
}'"#;
const POSIX_DUMP_ENV_SCRIPT: &str = "/usr/bin/env -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
awk 'BEGIN {
for (key in ENVIRON) {
printf "%s=%s%c", key, ENVIRON[key], 0
}
}'"#;
/usr/bin/env -0"#;
const POWERSHELL_DUMP_ENV_SCRIPT: &str = r#"$items = [ordered]@{}
Get-ChildItem Env: | Sort-Object Name | ForEach-Object {

View File

@@ -55,6 +55,7 @@ unset REMOVED_BY_HOOK
env_file
.capture_exports(&test_shell(), cwd.as_path(), &base_env)
.await?;
assert_eq!(std::fs::read(env_file.path())?, Vec::<u8>::new());
let mut env = base_env;
insert_env_file_paths(&env_file, &mut env);
@@ -138,7 +139,13 @@ export EXPLICIT_OVERRIDE='from-hook'
EnvironmentVariablePattern::new_case_insensitive("ALLOWED_*"),
EnvironmentVariablePattern::new_case_insensitive("EXPLICIT_*"),
],
r#set: HashMap::from([("EXPLICIT_OVERRIDE".to_string(), "from-policy".to_string())]),
r#set: HashMap::from([
("EXPLICIT_OVERRIDE".to_string(), "from-policy".to_string()),
(
"NOT_INCLUDED_OVERRIDE".to_string(),
"not-included".to_string(),
),
]),
..Default::default()
};
@@ -155,6 +162,55 @@ export EXPLICIT_OVERRIDE='from-hook'
Ok(())
}
#[cfg(not(windows))]
#[tokio::test]
async fn shell_env_file_capture_does_not_require_tools_on_configured_path() -> Result<()> {
let env_file = ShellEnvFile::new(ThreadId::new(), ShellEnvCapture::Posix)?;
let base_env = HashMap::from([("PATH".to_string(), "/missing".to_string())]);
std::fs::write(
env_file.path(),
"export CAPTURED_WITH_RESTRICTED_PATH='yes'\n",
)?;
let cwd = std::env::current_dir()?;
env_file
.capture_exports(&test_shell(), cwd.as_path(), &base_env)
.await?;
let mut env = base_env;
env_file.apply_exports(&mut env, &ShellEnvironmentPolicy::default());
assert_eq!(
env,
HashMap::from([
("PATH".to_string(), "/missing".to_string()),
(
"CAPTURED_WITH_RESTRICTED_PATH".to_string(),
"yes".to_string(),
),
])
);
Ok(())
}
#[cfg(not(windows))]
#[tokio::test]
async fn shell_env_file_reset_clears_previous_generation() -> Result<()> {
let env_file = ShellEnvFile::new(ThreadId::new(), ShellEnvCapture::Posix)?;
let base_env = HashMap::new();
std::fs::write(env_file.path(), "export PREVIOUS_GENERATION='stale'\n")?;
let cwd = std::env::current_dir()?;
env_file
.capture_exports(&test_shell(), cwd.as_path(), &base_env)
.await?;
env_file.reset_exports()?;
let mut env = base_env;
env_file.apply_exports(&mut env, &ShellEnvironmentPolicy::default());
assert_eq!(env, HashMap::new());
Ok(())
}
#[cfg(not(windows))]
fn test_shell() -> Shell {
Shell {

View File

@@ -15,6 +15,7 @@ use crate::exec::ExecCapturePolicy;
use crate::exec::StdoutStream;
use crate::exec::execute_exec_request;
use crate::exec_env::create_env;
use crate::hook_runtime::run_pending_session_start_hooks;
use crate::sandboxing::ExecRequest;
use crate::session::TurnInput;
use crate::session::turn_context::TurnContext;
@@ -120,6 +121,10 @@ pub(crate) async fn execute_user_shell_command(
session.send_event(turn_context.as_ref(), event).await;
}
if run_pending_session_start_hooks(&session, &turn_context).await {
return;
}
// Execute the user's script under their default shell when known; this
// allows commands that use shell features (pipes, &&, redirects, etc.).
// We do not source rc files or otherwise reformat the script.

View File

@@ -2767,6 +2767,40 @@ 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 session_start_env_file_exports_reach_first_user_shell_command() -> Result<()> {
skip_if_no_network!(Ok(()));
skip_if_windows!(Ok(()));
let server = start_mock_server().await;
let mut builder = test_codex()
.with_pre_build_hook(|home| {
if let Err(error) = write_session_start_hook_exporting_env(home) {
panic!("failed to write session env hook fixture: {error}");
}
})
.with_config(trust_discovered_hooks);
let test = builder.build(&server).await?;
test.codex
.submit(Op::RunUserShellCommand {
command: "printf '%s|%s|%s' \"$CODEX_SESSION_START_TEST\" \"${CODEX_ENV_FILE:+configured}\" \"${CLAUDE_ENV_FILE:+configured}\"".to_string(),
})
.await?;
let event = wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::ExecCommandEnd(_))
})
.await;
let EventMsg::ExecCommandEnd(event) = event else {
unreachable!();
};
assert_eq!(event.exit_code, 0);
assert_eq!(event.stdout, "from-session-start||");
Ok(())
}
async fn assert_pre_tool_use_rewrites_bash_surface(surface: BashRewriteSurface) -> Result<()> {
skip_if_no_network!(Ok(()));

View File

@@ -614,6 +614,32 @@ printf 'export FIRST=1\n' >> "$CLAUDE_ENV_FILE""#
);
}
#[cfg(not(windows))]
#[tokio::test]
async fn session_start_hooks_replace_previous_env_file_generation() {
use std::fs;
use tempfile::tempdir;
let env_file = tempfile::NamedTempFile::new().expect("create env file");
fs::write(env_file.path(), "export STALE=1\n").expect("seed stale env file");
let mut hook = handler();
hook.command = r#"printf 'export FRESH=1\n' >> "$CODEX_ENV_FILE""#.to_string();
let cwd = tempdir().expect("create cwd");
run_session_start_env_file_hooks(&[hook.clone()], &env_file, cwd.path()).await;
assert_eq!(
fs::read_to_string(env_file.path()).expect("read env file"),
"export FRESH=1\n"
);
run_session_start_env_file_hooks(&[hook], &env_file, cwd.path()).await;
assert_eq!(
fs::read_to_string(env_file.path()).expect("read env file"),
"export FRESH=1\n"
);
}
#[cfg(not(windows))]
async fn run_session_start_env_file_hooks(
handlers: &[ConfiguredHandler],

View File

@@ -1,4 +1,3 @@
use std::io::ErrorKind;
use std::path::Path;
use anyhow::Context;
@@ -38,6 +37,28 @@ pub(super) async fn execute_handlers<T>(
};
let env_file_path = Path::new(env_file_path);
let scratch_dir = env_file_path.parent().unwrap_or_else(|| Path::new("."));
if let Err(error) = fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(env_file_path)
.await
.with_context(|| format!("failed to reset {}", env_file_path.display()))
{
return handlers
.into_iter()
.enumerate()
.map(|(completion_order, handler)| {
let mut parsed = parse(
&handler,
error_result(anyhow::anyhow!("{error:#}")),
turn_id.clone(),
);
parsed.completion_order = completion_order;
parsed
})
.collect();
}
let mut pending = FuturesUnordered::new();
for (configured_order, handler) in handlers.into_iter().enumerate() {
@@ -92,17 +113,7 @@ pub(super) async fn execute_handlers<T>(
}
completed.sort_by_key(|execution| execution.configured_order);
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
}
};
let mut env_file_ends_with_newline = true;
// Merge each successful handler's isolated file in configured order. The
// core session layer will source this single canonical file after hooks

View File

@@ -109,7 +109,12 @@ where
pub fn inherited_env_var_allowed_by_policy(name: &str, policy: &ShellEnvironmentPolicy) -> bool {
(policy.ignore_default_excludes || !matches_any_pattern(name, DEFAULT_EXCLUDES.as_slice()))
&& !matches_any_pattern(name, &policy.exclude)
&& (policy.include_only.is_empty() || matches_any_pattern(name, &policy.include_only))
&& env_var_allowed_by_include_only(name, policy)
}
/// Returns whether a variable is permitted by the policy's allowlist.
pub fn env_var_allowed_by_include_only(name: &str, policy: &ShellEnvironmentPolicy) -> bool {
policy.include_only.is_empty() || matches_any_pattern(name, &policy.include_only)
}
fn matches_any_pattern(name: &str, patterns: &[EnvironmentVariablePattern]) -> bool {