From 027bd73a4c68a94c31fae6ed4356208a43c9066c Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Sat, 30 May 2026 14:29:00 -0700 Subject: [PATCH] Run SessionStart env-file hooks in parallel --- codex-rs/hooks/src/events/mod.rs | 1 + codex-rs/hooks/src/events/session_start.rs | 125 +++++++++-- .../src/events/session_start_env_file.rs | 212 ++++++++++++++++++ 3 files changed, 315 insertions(+), 23 deletions(-) create mode 100644 codex-rs/hooks/src/events/session_start_env_file.rs diff --git a/codex-rs/hooks/src/events/mod.rs b/codex-rs/hooks/src/events/mod.rs index 5ec24462b9..5ed8a7e9dc 100644 --- a/codex-rs/hooks/src/events/mod.rs +++ b/codex-rs/hooks/src/events/mod.rs @@ -4,5 +4,6 @@ pub mod permission_request; pub mod post_tool_use; pub mod pre_tool_use; pub mod session_start; +mod session_start_env_file; pub mod stop; pub mod user_prompt_submit; diff --git a/codex-rs/hooks/src/events/session_start.rs b/codex-rs/hooks/src/events/session_start.rs index 98c9df878f..4de3b2dc66 100644 --- a/codex-rs/hooks/src/events/session_start.rs +++ b/codex-rs/hooks/src/events/session_start.rs @@ -7,15 +7,13 @@ use codex_protocol::protocol::HookOutputEntry; use codex_protocol::protocol::HookOutputEntryKind; use codex_protocol::protocol::HookRunStatus; use codex_protocol::protocol::HookRunSummary; -use codex_protocol::shell_environment::CLAUDE_ENV_FILE_ENV_VAR; -use codex_protocol::shell_environment::CODEX_ENV_FILE_ENV_VAR; use codex_utils_absolute_path::AbsolutePathBuf; use super::common; +use super::session_start_env_file; use crate::engine::CommandShell; use crate::engine::ConfiguredHandler; use crate::engine::command_runner::CommandRunResult; -use crate::engine::command_runner::run_command; use crate::engine::dispatcher; use crate::engine::output_parser; use crate::schema::NullableString; @@ -181,21 +179,16 @@ pub(crate) async fn run( } }; - let results = if event_name == HookEventName::SessionStart - && (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()); - for (completion_order, handler) in matched.into_iter().enumerate() { - let result = run_command(shell, &handler, &input_json, request.cwd.as_path()).await; - let mut parsed = parse_completed(&handler, result, turn_id.clone()); - parsed.completion_order = completion_order; - results.push(parsed); - } - results + let results = if event_name == HookEventName::SessionStart { + session_start_env_file::execute_handlers( + shell, + matched, + input_json, + request.cwd.as_path(), + turn_id, + parse_completed, + ) + .await } else { dispatcher::execute_handlers( shell, @@ -528,14 +521,15 @@ mod tests { #[cfg(not(windows))] #[tokio::test] - async fn session_start_hooks_write_env_file_in_configured_order() { + async fn session_start_hooks_merge_env_file_in_configured_order_after_parallel_runs() { 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::CODEX_ENV_FILE_ENV_VAR; use super::SessionStartRequest; use super::SessionStartSource; use super::StartHookTarget; @@ -551,10 +545,25 @@ mod tests { env_file.path().display().to_string(), )]), }; + let cwd = tempdir().expect("create cwd"); + let marker = cwd.path().join("second-started"); let mut first = handler(); - first.command = r#"sleep 1; printf 'export FIRST=1\n' >> "$CODEX_ENV_FILE""#.to_string(); + first + .env + .insert("MARKER".to_string(), marker.display().to_string()); + first.command = r#"for attempt in 1 2 3 4 5 6 7 8 9 10; do + [ -f "$MARKER" ] && break + sleep 0.1 +done +[ -f "$MARKER" ] || exit 7 +printf 'export FIRST=1\n' >> "$CODEX_ENV_FILE""# + .to_string(); let mut second = handler(); - second.command = r#"grep -q 'export FIRST=1' "$CODEX_ENV_FILE" && printf 'export SECOND=1\n' >> "$CODEX_ENV_FILE""#.to_string(); + second + .env + .insert("MARKER".to_string(), marker.display().to_string()); + second.command = + r#"printf 'export SECOND=1\n' >> "$CODEX_ENV_FILE"; touch "$MARKER""#.to_string(); second.display_order = 1; let outcome = run( @@ -562,7 +571,7 @@ mod tests { &shell, SessionStartRequest { session_id: ThreadId::new(), - cwd: test_path_buf("/tmp").abs(), + cwd: cwd.path().to_path_buf().abs(), transcript_path: None, model: "gpt-test".to_string(), permission_mode: "default".to_string(), @@ -588,6 +597,76 @@ mod tests { ); } + #[cfg(not(windows))] + #[tokio::test] + async fn session_start_hooks_receive_matching_codex_and_claude_scratch_paths() { + use std::collections::HashMap; + use std::fs; + + use codex_protocol::ThreadId; + use codex_protocol::shell_environment::CLAUDE_ENV_FILE_ENV_VAR; + 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(), + ), + ( + CLAUDE_ENV_FILE_ENV_VAR.to_string(), + env_file.path().display().to_string(), + ), + ]), + }; + let mut hook = handler(); + hook.command = r#"[ "$CODEX_ENV_FILE" = "$CLAUDE_ENV_FILE" ] || exit 9 +printf 'export ALIAS_OK=1\n' >> "$CLAUDE_ENV_FILE""# + .to_string(); + let cwd = tempdir().expect("create cwd"); + + let outcome = run( + &[hook], + &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![HookRunStatus::Completed] + ); + assert_eq!( + fs::read_to_string(env_file.path()).expect("read env file"), + "export ALIAS_OK=1\n" + ); + } + fn handler() -> ConfiguredHandler { handler_for(HookEventName::SessionStart) } diff --git a/codex-rs/hooks/src/events/session_start_env_file.rs b/codex-rs/hooks/src/events/session_start_env_file.rs new file mode 100644 index 0000000000..2e3f64fed4 --- /dev/null +++ b/codex-rs/hooks/src/events/session_start_env_file.rs @@ -0,0 +1,212 @@ +use std::io::ErrorKind; +use std::path::Path; +use std::path::PathBuf; + +use anyhow::Context; +use anyhow::Result; +use futures::StreamExt; +use futures::stream::FuturesUnordered; +use tokio::fs; +use tokio::io::AsyncWriteExt; +use uuid::Uuid; + +use codex_protocol::shell_environment::CLAUDE_ENV_FILE_ENV_VAR; +use codex_protocol::shell_environment::CODEX_ENV_FILE_ENV_VAR; + +use crate::engine::CommandShell; +use crate::engine::ConfiguredHandler; +use crate::engine::command_runner::CommandRunResult; +use crate::engine::command_runner::run_command; +use crate::engine::dispatcher; + +/// Runs `SessionStart` handlers in parallel without sharing a writable env file. +/// +/// Each handler receives its own scratch `CODEX_ENV_FILE` / `CLAUDE_ENV_FILE` +/// path. After all handlers finish, their scratch-file contents are merged back +/// into the canonical env file in configured handler order. Core can then source +/// and capture that one canonical file exactly as it does for any other +/// `SessionStart` run. +pub(super) async fn execute_handlers( + shell: &CommandShell, + handlers: Vec, + input_json: String, + cwd: &Path, + turn_id: Option, + parse: fn(&ConfiguredHandler, CommandRunResult, Option) -> dispatcher::ParsedHandler, +) -> Vec> { + // No env file means there is nothing special to isolate; use the standard + // parallel dispatcher. + let Some(env_file_path) = shell + .session_start_env + .get(CODEX_ENV_FILE_ENV_VAR) + .or_else(|| shell.session_start_env.get(CLAUDE_ENV_FILE_ENV_VAR)) + .map(PathBuf::from) + else { + return dispatcher::execute_handlers(shell, handlers, input_json, cwd, turn_id, parse) + .await; + }; + + let file_name = env_file_path + .file_name() + .map(|name| name.to_string_lossy()) + .unwrap_or_else(|| "env".into()); + + let mut pending = FuturesUnordered::new(); + for (configured_order, handler) in handlers.into_iter().enumerate() { + // Each handler writes to an isolated env file, but still sees both + // compatibility aliases when the original shell had them. + let scratch_path = env_file_path.with_file_name(format!( + ".{file_name}.{configured_order}.{}.tmp", + 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 scratch_shell = shell.clone(); + let scratch_path_env = scratch_path.to_string_lossy().to_string(); + if scratch_shell + .session_start_env + .contains_key(CODEX_ENV_FILE_ENV_VAR) + { + scratch_shell + .session_start_env + .insert(CODEX_ENV_FILE_ENV_VAR.to_string(), scratch_path_env.clone()); + } + if scratch_shell + .session_start_env + .contains_key(CLAUDE_ENV_FILE_ENV_VAR) + { + scratch_shell + .session_start_env + .insert(CLAUDE_ENV_FILE_ENV_VAR.to_string(), scratch_path_env); + } + + Ok((scratch_path, scratch_shell)) + } + .await; + + let input_json = input_json.clone(); + let turn_id = turn_id.clone(); + pending.push(async move { + let (scratch_path, scratch_shell) = match scratch_result { + Ok(scratch) => scratch, + Err(error) => { + return HandlerExecution { + configured_order, + completion_order: 0, + handler, + run_result: error_result(error), + scratch_path: None, + turn_id, + }; + } + }; + let run_result = run_command(&scratch_shell, &handler, &input_json, cwd).await; + HandlerExecution { + configured_order, + completion_order: 0, + handler, + run_result, + scratch_path: Some(scratch_path), + turn_id, + } + }); + } + + let mut completed = Vec::new(); + let mut completion_order = 0; + while let Some(mut execution) = pending.next().await { + execution.completion_order = completion_order; + completion_order += 1; + completed.push(execution); + } + 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. + for execution in &mut completed { + let Some(scratch_path) = execution.scratch_path.take() else { + continue; + }; + let result: Result<()> = async { + let contents = fs::read(&scratch_path).await.with_context(|| { + format!( + "failed to read SessionStart scratch env file {}", + scratch_path.display() + ) + })?; + if contents.is_empty() { + return Ok(()); + } + + let mut file = fs::OpenOptions::new() + .append(true) + .create(true) + .open(&env_file_path) + .await + .with_context(|| format!("failed to open {}", 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(()) + } + .await; + + let cleanup_result = fs::remove_file(&scratch_path).await; + if let Err(error) = cleanup_result + && error.kind() != ErrorKind::NotFound + { + tracing::warn!( + "failed to remove SessionStart scratch env file {}: {error}", + scratch_path.display() + ); + } + if let Err(error) = result { + execution.run_result = error_result(error); + } + } + + completed + .into_iter() + .map(|execution| { + let mut parsed = parse(&execution.handler, execution.run_result, execution.turn_id); + parsed.completion_order = execution.completion_order; + parsed + }) + .collect() +} + +struct HandlerExecution { + configured_order: usize, + completion_order: usize, + handler: ConfiguredHandler, + run_result: CommandRunResult, + scratch_path: Option, + turn_id: Option, +} + +fn error_result(error: anyhow::Error) -> CommandRunResult { + let timestamp = chrono::Utc::now().timestamp(); + CommandRunResult { + started_at: timestamp, + completed_at: timestamp, + duration_ms: 0, + exit_code: None, + stdout: String::new(), + stderr: String::new(), + error: Some(format!("{error:#}")), + } +}