fix(oh-my-zsh): fix remote oh-my-zsh issue (#496)

* fix(oh-my-zsh): fix remote oh-my-zsh issue

src/execution_context.rs

@@ -5,6 +5,7 @@ use crate::sudo::Sudo;
 use crate::utils::{require_option, REQUIRE_SUDO};
 use crate::{config::Config, executor::Executor};
 use color_eyre::eyre::Result;
+use std::env::var;
 use std::path::Path;
 use std::sync::Mutex;
 
@@ -17,16 +18,20 @@ pub struct ExecutionContext<'a> {
     /// This is used in `./steps/remote/ssh.rs`, where we want to run `topgrade` in a new
     /// tmux window for each remote.
     tmux_session: Mutex<Option<String>>,
+    /// True if topgrade is running under ssh.
+    under_ssh: bool,
 }
 
 impl<'a> ExecutionContext<'a> {
     pub fn new(run_type: RunType, sudo: Option<Sudo>, git: &'a Git, config: &'a Config) -> Self {
+        let under_ssh = var("SSH_CLIENT").is_ok() || var("SSH_TTY").is_ok();
         Self {
             run_type,
             sudo,
             git,
             config,
             tmux_session: Mutex::new(None),
+            under_ssh,
         }
     }
 
@@ -51,6 +56,10 @@ impl<'a> ExecutionContext<'a> {
         self.config
     }
 
+    pub fn under_ssh(&self) -> bool {
+        self.under_ssh
+    }
+
     pub fn set_tmux_session(&self, session_name: String) {
         self.tmux_session.lock().unwrap().replace(session_name);
     }

src/steps/zsh.rs

@@ -165,6 +165,28 @@ pub fn run_zim(ctx: &ExecutionContext) -> Result<()> {
 
 pub fn run_oh_my_zsh(ctx: &ExecutionContext) -> Result<()> {
     require("zsh")?;
+
+    // When updating `oh-my-zsh` on a remote machine through topgrade, the
+    // following processes will be created:
+    //
+    // SSH -> ZSH -> ZSH ($SHELL) -> topgrade -> ZSH
+    //
+    // The first ZSH process, won't source zshrc (as it is a login shell),
+    // and thus it won't have the ZSH environment variable, as a result, the
+    // children processes won't get it either, so we source the zshrc and set
+    // the ZSH variable for topgrade here.
+    if ctx.under_ssh() {
+        let zshrc_path = zshrc().require()?;
+        let output = Command::new("zsh")
+            .args([
+                "-c",
+                // ` > /dev/null` is used in case the user's zshrc will have some stdout output.
+                format!("source {} > /dev/null && echo $ZSH", zshrc_path.display()).as_str(),
+            ])
+            .output_checked_utf8()?;
+        env::set_var("ZSH", output.stdout.trim());
+    }
+
     let oh_my_zsh = env::var("ZSH")
         .map(PathBuf::from)
         // default to `~/.oh-my-zsh`