Run SessionStart env-file hooks in parallel

2026-10-21 17:48:44 +00:00 · 2026-05-30 14:29:00 -07:00
parent 1f0fe5b8d9
commit 027bd73a4c
3 changed files with 315 additions and 23 deletions
@@ -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;
@@ -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<_>>(),
+            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)
    }
@@ -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<T>(
+    shell: &CommandShell,
+    handlers: Vec<ConfiguredHandler>,
+    input_json: String,
+    cwd: &Path,
+    turn_id: Option<String>,
+    parse: fn(&ConfiguredHandler, CommandRunResult, Option<String>) -> dispatcher::ParsedHandler<T>,
+) -> Vec<dispatcher::ParsedHandler<T>> {
+    // 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<PathBuf>,
+    turn_id: Option<String>,
+}
+
+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:#}")),
+    }
+}