From 2f2317b70b8056a826a6172583e38ebecca3fa78 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Sat, 30 May 2026 21:27:27 -0700 Subject: [PATCH] Simplify SessionStart environment overlays --- codex-rs/core/src/session/mod.rs | 9 --- codex-rs/core/src/session_start_env.rs | 74 +++---------------- codex-rs/core/src/tasks/user_shell.rs | 2 +- codex-rs/core/src/tools/runtimes/shell.rs | 2 +- .../core/src/tools/runtimes/unified_exec.rs | 2 +- codex-rs/hooks/src/engine/output_parser.rs | 33 +++------ codex-rs/hooks/src/events/session_start.rs | 23 +++--- 7 files changed, 33 insertions(+), 112 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index e01924bb29..3f3ff36cee 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3230,15 +3230,6 @@ impl Session { self.services.session_start_env.apply(env); } - pub(crate) fn extend_session_start_env_snapshot_overrides( - &self, - overrides: &mut HashMap, - ) { - self.services - .session_start_env - .extend_snapshot_overrides(overrides); - } - pub(crate) async fn current_rollout_path(&self) -> anyhow::Result> { let Some(live_thread) = self.live_thread() else { return Ok(None); diff --git a/codex-rs/core/src/session_start_env.rs b/codex-rs/core/src/session_start_env.rs index c9d9ee0c01..3c5139424f 100644 --- a/codex-rs/core/src/session_start_env.rs +++ b/codex-rs/core/src/session_start_env.rs @@ -22,59 +22,24 @@ impl SessionStartEnvOverlay { .unwrap_or_else(PoisonError::into_inner) .iter() { - if is_runtime_owned(key) { + if key == CODEX_THREAD_ID_ENV_VAR + || cfg!(windows) && key.eq_ignore_ascii_case(CODEX_THREAD_ID_ENV_VAR) + { continue; } - insert_env_var(env, key.clone(), value.clone()); - } - } - - pub(crate) fn extend_snapshot_overrides(&self, overrides: &mut HashMap) { - for (key, value) in self - .values - .read() - .unwrap_or_else(PoisonError::into_inner) - .iter() - { - if is_runtime_owned(key) { - continue; + #[cfg(windows)] + if let Some(existing) = env + .keys() + .find(|candidate| candidate.eq_ignore_ascii_case(key)) + .cloned() + { + env.remove(&existing); } - insert_env_var(overrides, key.clone(), value.clone()); + env.insert(key.clone(), value.clone()); } } } -fn is_runtime_owned(key: &str) -> bool { - env_key_eq(key, CODEX_THREAD_ID_ENV_VAR) -} - -fn insert_env_var(env: &mut HashMap, key: String, value: String) { - remove_env_var(env, &key); - env.insert(key, value); -} - -fn remove_env_var(env: &mut HashMap, key: &str) { - if let Some(existing) = env - .keys() - .find(|candidate| env_key_eq(candidate, key)) - .cloned() - { - env.remove(&existing); - } -} - -fn env_key_eq(candidate: &str, key: &str) -> bool { - #[cfg(windows)] - { - candidate.eq_ignore_ascii_case(key) - } - - #[cfg(not(windows))] - { - candidate == key - } -} - #[cfg(test)] mod tests { use pretty_assertions::assert_eq; @@ -106,21 +71,4 @@ mod tests { ]) ); } - - #[test] - fn snapshot_overrides_include_hook_env() { - let overlay = SessionStartEnvOverlay::default(); - overlay.replace(HashMap::from([( - "PLUGIN_HOME".to_string(), - "/plugin".to_string(), - )])); - let mut overrides = HashMap::new(); - - overlay.extend_snapshot_overrides(&mut overrides); - - assert_eq!( - overrides, - HashMap::from([("PLUGIN_HOME".to_string(), "/plugin".to_string())]) - ); - } } diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index ef73ca5dde..f755d3caca 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -152,7 +152,7 @@ pub(crate) async fn execute_user_shell_command( } } let mut snapshot_overrides = turn_context.shell_environment_policy.r#set.clone(); - session.extend_session_start_env_snapshot_overrides(&mut snapshot_overrides); + session.apply_session_start_env(&mut snapshot_overrides); let exec_command = maybe_wrap_shell_lc_with_snapshot( &display_command, session_shell.as_ref(), diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index d286b9272b..a1890cd29f 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -261,7 +261,7 @@ impl ToolRuntime for ShellRuntime { }; let mut explicit_env_overrides = explicit_env_overrides; ctx.session - .extend_session_start_env_snapshot_overrides(&mut explicit_env_overrides); + .apply_session_start_env(&mut explicit_env_overrides); let command = maybe_wrap_shell_lc_with_snapshot( &req.command, session_shell.as_ref(), diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index 996f840a60..fb83bd5a39 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -295,7 +295,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt let mut explicit_env_overrides = explicit_env_overrides; if !environment_is_remote { ctx.session - .extend_session_start_env_snapshot_overrides(&mut explicit_env_overrides); + .apply_session_start_env(&mut explicit_env_overrides); } let command = if environment_is_remote { base_command.to_vec() diff --git a/codex-rs/hooks/src/engine/output_parser.rs b/codex-rs/hooks/src/engine/output_parser.rs index efb9bd6875..6ec56c78f9 100644 --- a/codex-rs/hooks/src/engine/output_parser.rs +++ b/codex-rs/hooks/src/engine/output_parser.rs @@ -101,37 +101,24 @@ pub(crate) fn parse_session_start(stdout: &str) -> Option { .map(|output| (output.additional_context, output.env)) .unwrap_or_default(); let invalid_reason = invalid_session_start_env_reason(&env); - Some(session_start_output( - wire.universal, + Some(SessionStartOutput { + universal: UniversalOutput::from(wire.universal), additional_context, env, invalid_reason, - )) + }) } pub(crate) fn parse_subagent_start(stdout: &str) -> Option { let wire: SubagentStartCommandOutputWire = parse_json(stdout)?; - Some(session_start_output( - wire.universal, - wire.hook_specific_output + Some(SessionStartOutput { + universal: UniversalOutput::from(wire.universal), + additional_context: wire + .hook_specific_output .and_then(|output| output.additional_context), - HashMap::new(), - None, - )) -} - -fn session_start_output( - universal: HookUniversalOutputWire, - additional_context: Option, - env: HashMap, - invalid_reason: Option, -) -> SessionStartOutput { - SessionStartOutput { - universal: UniversalOutput::from(universal), - additional_context, - env, - invalid_reason, - } + env: HashMap::new(), + invalid_reason: None, + }) } fn invalid_session_start_env_reason(env: &HashMap) -> Option { diff --git a/codex-rs/hooks/src/events/session_start.rs b/codex-rs/hooks/src/events/session_start.rs index 857693b17b..159b259994 100644 --- a/codex-rs/hooks/src/events/session_start.rs +++ b/codex-rs/hooks/src/events/session_start.rs @@ -376,25 +376,20 @@ fn merge_session_start_env<'a>( let mut merged = HashMap::new(); for env in envs { for (key, value) in env { - insert_session_start_env(&mut merged, key.clone(), value.clone()); + #[cfg(windows)] + if let Some(existing) = merged + .keys() + .find(|candidate| candidate.eq_ignore_ascii_case(key)) + .cloned() + { + merged.remove(&existing); + } + merged.insert(key.clone(), value.clone()); } } merged } -fn insert_session_start_env(env: &mut HashMap, key: String, value: String) { - #[cfg(windows)] - if let Some(existing) = env - .keys() - .find(|candidate| candidate.eq_ignore_ascii_case(&key)) - .cloned() - { - env.remove(&existing); - } - - env.insert(key, value); -} - #[cfg(test)] mod tests { use std::collections::HashMap;