Run hooks with the captured session environment (#39314)

## What changed

- Capture the process environment when the hook registry is created and reuse
  that snapshot across configuration reloads.
- Clear the live environment before launching command hooks and legacy notify
  commands, then apply hook-specific overrides and scrub non-inheritable
  credentials.
- Resolve the default shell from the captured environment.

## Testing

Add coverage for snapshot replay, overrides, credential scrubbing, default
shell selection, non-Unicode values, and runtime reconfiguration.

GitOrigin-RevId: fee60c88e842980cdfc1bd49b14b62b9a56b08cd
This commit is contained in:
Abhinav
2026-08-18 23:17:34 +00:00
committed by copyberry
parent fb356f3d2c
commit 4d8c664a49
6 changed files with 241 additions and 32 deletions

View File

@@ -1,4 +1,7 @@
use std::collections::HashMap;
#[cfg(not(windows))]
use std::ffi::OsStr;
use std::ffi::OsString;
use std::io::ErrorKind;
use std::path::Path;
use std::process::Stdio;
@@ -44,6 +47,7 @@ const MAX_CONCURRENT_ASYNC_HOOKS: usize = 8;
#[derive(Clone)]
pub(crate) struct CommandHookRuntime {
shell: CommandShell,
environment: Arc<Vec<(OsString, OsString)>>,
result_sender: Sender<HookCompletedEvent>,
state: Arc<Mutex<CommandHookRuntimeState>>,
output_spiller: HookOutputSpiller,
@@ -66,11 +70,13 @@ impl Default for CommandHookRuntimeState {
impl CommandHookRuntime {
pub(crate) fn new(
shell: CommandShell,
environment: Arc<Vec<(OsString, OsString)>>,
thread_id: ThreadId,
result_sender: Sender<HookCompletedEvent>,
) -> Self {
Self {
shell,
environment,
result_sender,
state: Arc::new(Mutex::new(CommandHookRuntimeState::default())),
output_spiller: HookOutputSpiller::new(thread_id),
@@ -86,6 +92,7 @@ impl CommandHookRuntime {
pub(crate) fn reconfigured(&self, shell: CommandShell) -> Self {
Self {
shell,
environment: Arc::clone(&self.environment),
result_sender: self.result_sender.clone(),
state: Arc::clone(&self.state),
output_spiller: self.output_spiller.clone(),
@@ -195,7 +202,7 @@ pub(crate) async fn run_command(
let started_at = chrono::Utc::now().timestamp();
let started = Instant::now();
let mut command = build_command(&runtime.shell, command, env);
let mut command = build_command(&runtime.shell, command, &runtime.environment, env);
command
.current_dir(cwd)
.stdin(Stdio::piped())
@@ -372,10 +379,11 @@ fn finish_command_run(
fn build_command(
shell: &CommandShell,
command_line: &str,
environment: &[(OsString, OsString)],
env: &HashMap<String, String>,
) -> Command {
let mut command = if shell.program.is_empty() {
default_shell_command()
default_shell_command(environment)
} else {
Command::new(&shell.program)
};
@@ -398,27 +406,41 @@ fn build_command(
#[cfg(not(windows))]
command.arg(command_line);
}
// Replay the session snapshot instead of inheriting the live process environment.
command.env_clear();
command.envs(environment.iter().cloned());
command.envs(env);
scrub_non_inheritable_env_vars(command.as_std_mut());
command
}
fn default_shell_command() -> Command {
fn default_shell_command(environment: &[(OsString, OsString)]) -> Command {
#[cfg(windows)]
{
let comspec = std::env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".to_string());
let mut command = Command::new(comspec);
command.arg("/C");
command
}
let (environment_variable, fallback_program, argument) = ("COMSPEC", "cmd.exe", "/C");
#[cfg(not(windows))]
{
let shell = std::env::var("SHELL").unwrap_or_else(|_| "/bin/sh".to_string());
let mut command = Command::new(shell);
command.arg("-lc");
command
}
let (environment_variable, fallback_program, argument) = ("SHELL", "/bin/sh", "-lc");
let program = environment
.iter()
.find(|(key, _)| {
#[cfg(windows)]
{
key.to_str()
.is_some_and(|key| key.eq_ignore_ascii_case(environment_variable))
}
#[cfg(not(windows))]
{
key == OsStr::new(environment_variable)
}
})
.map(|(_, value)| value.clone())
.unwrap_or_else(|| OsString::from(fallback_program));
let mut command = Command::new(program);
command.arg(argument);
command
}
#[cfg(test)]

View File

@@ -1,7 +1,12 @@
use std::collections::HashMap;
use std::ffi::OsStr;
use std::ffi::OsString;
#[cfg(windows)]
use std::fs;
#[cfg(unix)]
use std::os::unix::ffi::OsStringExt;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
use async_channel::Receiver;
@@ -27,6 +32,7 @@ use super::CommandHookRuntime;
use super::CommandShell;
use super::ConfiguredHandler;
use super::MAX_CONCURRENT_ASYNC_HOOKS;
use super::build_command;
use super::run_command;
use crate::events::user_prompt_submit::UserPromptSubmitRequest;
@@ -74,7 +80,12 @@ async fn cmd_shell_runs_quoted_hook_command_path() {
for shell in shells {
let (result_sender, _result_receiver) = async_channel::unbounded();
let runtime = CommandHookRuntime::new(shell, ThreadId::new(), result_sender);
let runtime = CommandHookRuntime::new(
shell,
Arc::new(std::env::vars_os().collect()),
ThreadId::new(),
result_sender,
);
let result = run_command(&runtime, &handler, &command, &env, "{}", temp.path()).await;
assert_eq!(result.exit_code, Some(0), "stderr: {}", result.stderr);
@@ -157,9 +168,107 @@ async fn command_hook_does_not_expose_configured_noise_auth_token() {
assert_eq!(result.error, None);
}
#[test]
fn build_command_replays_snapshot_before_hook_overrides_and_scrubbing() {
#[cfg(unix)]
let non_unicode_value = OsString::from_vec(vec![b'v', 0xff]);
let environment = vec![
(
OsString::from("CODEX_HOOK_SNAPSHOT"),
OsString::from("captured"),
),
(
OsString::from("CODEX_HOOK_OVERRIDE"),
OsString::from("captured"),
),
(
OsString::from(CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR),
OsString::from("captured-noise-token"),
),
#[cfg(unix)]
(
OsString::from("CODEX_HOOK_NON_UNICODE"),
non_unicode_value.clone(),
),
];
let env = HashMap::from([
("CODEX_HOOK_OVERRIDE".to_string(), "configured".to_string()),
("CODEX_HOOK_SAFE_ENV".to_string(), "visible".to_string()),
(
CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR.to_string(),
"configured-noise-token".to_string(),
),
]);
let command = build_command(
&CommandShell {
program: "configured-shell".to_string(),
args: vec!["-c".to_string()],
},
"echo hook-ran",
&environment,
&env,
);
assert_eq!(
configured_environment_value(&command, "CODEX_HOOK_SNAPSHOT"),
Some(Some(OsString::from("captured")))
);
assert_eq!(
configured_environment_value(&command, "CODEX_HOOK_OVERRIDE"),
Some(Some(OsString::from("configured")))
);
assert_eq!(
configured_environment_value(&command, "CODEX_HOOK_SAFE_ENV"),
Some(Some(OsString::from("visible")))
);
assert_eq!(
configured_environment_value(&command, CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR),
None
);
#[cfg(unix)]
assert_eq!(
configured_environment_value(&command, "CODEX_HOOK_NON_UNICODE"),
Some(Some(non_unicode_value))
);
}
#[test]
fn fallback_shell_uses_snapshot() {
#[cfg(windows)]
let (name, program) = ("comspec", r"C:\captured\cmd.exe");
#[cfg(not(windows))]
let (name, program) = ("SHELL", "/captured/shell");
let command = build_command(
&CommandShell {
program: String::new(),
args: Vec::new(),
},
"echo hook-ran",
&[(OsString::from(name), OsString::from(program))],
&HashMap::new(),
);
assert_eq!(command.as_std().get_program(), OsStr::new(program));
#[cfg(not(windows))]
assert_eq!(
command
.as_std()
.get_args()
.map(OsStr::to_os_string)
.collect::<Vec<_>>(),
vec![OsString::from("-lc"), OsString::from("echo hook-ran")]
);
}
const ASYNC_HOOK_TEST_TIMEOUT: Duration = Duration::from_secs(30);
fn runtime() -> (CommandHookRuntime, Receiver<HookCompletedEvent>) {
runtime_with_environment(Arc::new(std::env::vars_os().collect()))
}
fn runtime_with_environment(
environment: Arc<Vec<(OsString, OsString)>>,
) -> (CommandHookRuntime, Receiver<HookCompletedEvent>) {
let thread_id = ThreadId::new();
let (result_sender, result_receiver) = async_channel::unbounded();
let runtime = CommandHookRuntime::new(
@@ -167,12 +276,24 @@ fn runtime() -> (CommandHookRuntime, Receiver<HookCompletedEvent>) {
program: String::new(),
args: Vec::new(),
},
environment,
thread_id,
result_sender,
);
(runtime, result_receiver)
}
fn configured_environment_value(
command: &tokio::process::Command,
name: &str,
) -> Option<Option<OsString>> {
command
.as_std()
.get_envs()
.find(|(key, _)| *key == OsStr::new(name))
.map(|(_, value)| value.map(OsStr::to_os_string))
}
fn write_handler(temp: &TempDir, source: &str) -> ConfiguredHandler {
let script_path = temp.path().join("async_hook.py");
std::fs::write(&script_path, source).expect("write async test hook");
@@ -250,12 +371,18 @@ print('{"systemMessage": 123}')
#[tokio::test]
async fn async_hook_result_survives_runtime_reconfiguration() {
let temp = TempDir::new().expect("async test directory");
let (previous, results) = runtime();
let mut environment = std::env::vars_os().collect::<Vec<_>>();
environment.push((
OsString::from("CODEX_HOOK_CAPTURED_ENV"),
OsString::from("captured"),
));
let (previous, results) = runtime_with_environment(Arc::new(environment));
let release_path = temp.path().join("release");
let handler = write_handler(
&temp,
&format!(
r#"import json
import os
from pathlib import Path
import sys
import time
@@ -263,7 +390,7 @@ import time
json.load(sys.stdin)
while not Path(r"{}").exists():
time.sleep(0.01)
print("survived reconfiguration")
print(os.environ["CODEX_HOOK_CAPTURED_ENV"])
"#,
release_path.display()
),
@@ -274,6 +401,10 @@ print("survived reconfiguration")
program: String::new(),
args: Vec::new(),
});
assert!(Arc::ptr_eq(
&previous.environment,
&reconfigured.environment
));
std::fs::write(release_path, "ready").expect("release async hook");
let hook_result = timeout(ASYNC_HOOK_TEST_TIMEOUT, results.recv())
@@ -284,7 +415,7 @@ print("survived reconfiguration")
hook_result.run.entries,
vec![HookOutputEntry {
kind: HookOutputEntryKind::Context,
text: "survived reconfiguration".to_string(),
text: "captured".to_string(),
}]
);

View File

@@ -51,7 +51,12 @@ fn cwd() -> AbsolutePathBuf {
fn command_runtime(shell: CommandShell) -> CommandHookRuntime {
let (result_sender, _result_receiver) = async_channel::unbounded();
CommandHookRuntime::new(shell, ThreadId::new(), result_sender)
CommandHookRuntime::new(
shell,
Arc::new(std::env::vars_os().collect()),
ThreadId::new(),
result_sender,
)
}
pub(crate) fn mcp_executor() -> Arc<dyn HookMcpExecutor> {

View File

@@ -1,3 +1,4 @@
use std::ffi::OsString;
use std::process::Stdio;
use std::sync::Arc;
@@ -7,7 +8,7 @@ use crate::Hook;
use crate::HookEvent;
use crate::HookPayload;
use crate::HookResult;
use crate::command_from_argv;
use crate::registry::command_from_argv;
/// Legacy notify payload appended as the final argv argument for backward compatibility.
#[derive(Debug, Clone, PartialEq, Serialize)]
@@ -40,14 +41,16 @@ pub fn legacy_notify_json(payload: &HookPayload) -> Result<String, serde_json::E
}
}
pub fn notify_hook(argv: Vec<String>) -> Hook {
// TODO: Remove this hook and its environment plumbing when legacy `notify` support is removed.
pub(crate) fn notify_hook(argv: Vec<String>, environment: Arc<Vec<(OsString, OsString)>>) -> Hook {
let argv = Arc::new(argv);
Hook {
name: "legacy_notify".to_string(),
func: Arc::new(move |payload: &HookPayload| {
let argv = Arc::clone(&argv);
let environment = Arc::clone(&environment);
Box::pin(async move {
let mut command = match command_from_argv(&argv) {
let mut command = match command_from_argv(&argv, environment.iter().cloned()) {
Some(command) => command,
None => return HookResult::Success,
};
@@ -73,6 +76,7 @@ pub fn notify_hook(argv: Vec<String>) -> Hook {
mod tests {
use anyhow::Result;
use codex_protocol::ThreadId;
use codex_protocol::shell_environment::CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR;
use codex_utils_absolute_path::test_support::PathBufExt;
use codex_utils_absolute_path::test_support::test_path_buf;
use pretty_assertions::assert_eq;
@@ -143,4 +147,37 @@ mod tests {
Ok(())
}
#[test]
fn legacy_notify_command_replays_session_snapshot_and_scrubs_credentials() {
let argv = vec!["notify-command".to_string()];
let environment = vec![
(
OsString::from("CODEX_LEGACY_NOTIFY_SNAPSHOT"),
OsString::from("captured"),
),
(
OsString::from(CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR),
OsString::from("restricted-token"),
),
];
let command = command_from_argv(&argv, environment)
.expect("legacy notification command should be configured");
let configured_environment = command
.as_std()
.get_envs()
.filter_map(|(name, value)| {
value.map(|value| (name.to_os_string(), value.to_os_string()))
})
.collect::<Vec<_>>();
assert_eq!(
configured_environment,
vec![(
OsString::from("CODEX_LEGACY_NOTIFY_SNAPSHOT"),
OsString::from("captured"),
)]
);
}
}

View File

@@ -74,13 +74,11 @@ pub use events::stop::StopRequest;
pub use events::user_prompt_submit::UserPromptSubmitOutcome;
pub use events::user_prompt_submit::UserPromptSubmitRequest;
pub use legacy_notify::legacy_notify_json;
pub use legacy_notify::notify_hook;
pub use mcp::HookMcpCall;
pub use mcp::HookMcpExecutor;
pub use registry::HookListOutcome;
pub use registry::Hooks;
pub use registry::HooksConfig;
pub use registry::command_from_argv;
pub use registry::list_hooks;
pub use schema::write_schema_fixtures;
pub use types::Hook;

View File

@@ -30,6 +30,7 @@ use codex_config::ConfigLayerStack;
use codex_plugin::PluginHookSource;
use codex_protocol::ThreadId;
use codex_protocol::shell_environment::scrub_non_inheritable_env_vars;
use std::ffi::OsString;
use std::sync::Arc;
use std::time::Duration;
use tokio::process::Command;
@@ -54,6 +55,9 @@ pub struct HookListOutcome {
#[derive(Clone)]
pub struct Hooks {
// TODO: Once legacy `notify` is removed, capture this snapshot in `CommandHookRuntime::new`
// and remove the environment plumbing from `Hooks` and `from_config`.
environment: Arc<Vec<(OsString, OsString)>>,
after_agent: Vec<Hook>,
engine: ClaudeHooksEngine,
}
@@ -67,8 +71,9 @@ impl Hooks {
mcp_executor: Arc<dyn HookMcpExecutor>,
) -> anyhow::Result<(Self, Receiver<codex_protocol::protocol::HookCompletedEvent>)> {
let (result_sender, result_receiver) = async_channel::unbounded();
let hooks = Self::from_config(config, mcp_executor, |shell| {
CommandHookRuntime::new(shell, thread_id, result_sender)
let environment = Arc::new(std::env::vars_os().collect());
let hooks = Self::from_config(config, mcp_executor, Arc::clone(&environment), |shell| {
CommandHookRuntime::new(shell, environment, thread_id, result_sender)
});
let required_load_errors = hooks.engine.required_load_errors();
if !required_load_errors.is_empty() {
@@ -82,20 +87,24 @@ impl Hooks {
/// Preserve in-flight background hooks while applying a refreshed configuration.
pub fn reconfigured(&self, config: HooksConfig) -> Self {
Self::from_config(config, Arc::clone(&self.engine.mcp_executor), |shell| {
self.engine.command_runtime.reconfigured(shell)
})
Self::from_config(
config,
Arc::clone(&self.engine.mcp_executor),
Arc::clone(&self.environment),
|shell| self.engine.command_runtime.reconfigured(shell),
)
}
fn from_config(
config: HooksConfig,
mcp_executor: Arc<dyn HookMcpExecutor>,
environment: Arc<Vec<(OsString, OsString)>>,
build_runtime: impl FnOnce(CommandShell) -> CommandHookRuntime,
) -> Self {
let after_agent = config
.legacy_notify_argv
.filter(|argv| !argv.is_empty() && !argv[0].is_empty())
.map(crate::notify_hook)
.map(|argv| crate::legacy_notify::notify_hook(argv, Arc::clone(&environment)))
.into_iter()
.collect();
let command_runtime = build_runtime(CommandShell {
@@ -112,6 +121,7 @@ impl Hooks {
mcp_executor,
);
Self {
environment,
after_agent,
engine,
}
@@ -278,13 +288,19 @@ pub fn list_hooks(config: HooksConfig) -> HookListOutcome {
}
}
pub fn command_from_argv(argv: &[String]) -> Option<Command> {
// TODO: Remove this legacy-notify-only command builder when `notify` support is removed.
pub(crate) fn command_from_argv(
argv: &[String],
environment: impl IntoIterator<Item = (OsString, OsString)>,
) -> Option<Command> {
let (program, args) = argv.split_first()?;
if program.is_empty() {
return None;
}
let mut command = Command::new(program);
command.args(args);
command.env_clear();
command.envs(environment);
scrub_non_inheritable_env_vars(command.as_std_mut());
Some(command)
}