diff --git a/crates/forge_main/src/banner.rs b/crates/forge_main/src/banner.rs index 1827ddaa52..6842fae8a7 100644 --- a/crates/forge_main/src/banner.rs +++ b/crates/forge_main/src/banner.rs @@ -105,14 +105,21 @@ pub fn display(cli_mode: bool) -> io::Result<()> { println!("{banner}\n"); - // Encourage zsh integration after the banner - if !cli_mode { + if !cli_mode && !is_zsh_plugin_loaded() { display_zsh_encouragement(); } Ok(()) } +/// Returns `true` when the zsh plugin exported `_FORGE_PLUGIN_LOADED` into the +/// environment inherited by this forge process. +fn is_zsh_plugin_loaded() -> bool { + std::env::var("_FORGE_PLUGIN_LOADED") + .map(|v| !v.is_empty()) + .unwrap_or(false) +} + /// Encourages users to use the zsh plugin for a better experience. fn display_zsh_encouragement() { let tip = DisplayBox::new(vec![ @@ -136,3 +143,65 @@ fn display_zsh_encouragement() { ]); println!("{}", tip); } + +#[cfg(test)] +mod tests { + use std::env; + use std::ffi::OsString; + use std::sync::{LazyLock, Mutex, MutexGuard}; + + use super::*; + + static ENV_MUTEX: LazyLock> = LazyLock::new(|| Mutex::new(())); + + struct EnvGuard { + key: &'static str, + original_value: Option, + _lock: MutexGuard<'static, ()>, + } + + impl EnvGuard { + fn set(key: &'static str, value: &str) -> Self { + let lock = ENV_MUTEX.lock().unwrap_or_else(|error| error.into_inner()); + let original_value = env::var_os(key); + unsafe { + env::set_var(key, value); + } + Self { key, original_value, _lock: lock } + } + + fn unset(key: &'static str) -> Self { + let lock = ENV_MUTEX.lock().unwrap_or_else(|error| error.into_inner()); + let original_value = env::var_os(key); + unsafe { + env::remove_var(key); + } + Self { key, original_value, _lock: lock } + } + } + + impl Drop for EnvGuard { + fn drop(&mut self) { + match &self.original_value { + Some(value) => unsafe { + env::set_var(self.key, value); + }, + None => unsafe { + env::remove_var(self.key); + }, + } + } + } + + #[test] + fn test_is_zsh_plugin_loaded_when_env_var_set() { + let _guard = EnvGuard::set("_FORGE_PLUGIN_LOADED", "1700000000"); + assert!(is_zsh_plugin_loaded()); + } + + #[test] + fn test_is_zsh_plugin_loaded_when_env_var_missing() { + let _guard = EnvGuard::unset("_FORGE_PLUGIN_LOADED"); + assert!(!is_zsh_plugin_loaded()); + } +} diff --git a/crates/forge_main/src/zsh/plugin.rs b/crates/forge_main/src/zsh/plugin.rs index 5e96121952..1dd04388b8 100644 --- a/crates/forge_main/src/zsh/plugin.rs +++ b/crates/forge_main/src/zsh/plugin.rs @@ -43,8 +43,8 @@ pub fn generate_zsh_plugin() -> Result { output.push_str("\n# --- Clap Completions ---\n"); output.push_str(&completions_str); - // Set environment variable to indicate plugin is loaded (with timestamp) - output.push_str("\n_FORGE_PLUGIN_LOADED=$(date +%s)\n"); + // Exported so child `forge` processes can detect an active plugin. + output.push_str("\nexport _FORGE_PLUGIN_LOADED=$(date +%s)\n"); Ok(output) } @@ -54,8 +54,8 @@ pub fn generate_zsh_theme() -> Result { let mut content = super::normalize_script(include_str!("../../../../shell-plugin/forge.theme.zsh")); - // Set environment variable to indicate theme is loaded (with timestamp) - content.push_str("\n_FORGE_THEME_LOADED=$(date +%s)\n"); + // Exported so child `forge` processes can detect an active theme. + content.push_str("\nexport _FORGE_THEME_LOADED=$(date +%s)\n"); Ok(content) }