End SSH session when closing tmux session

I have added the following to my .zshrc so that tmux will either create a new or attach to an existing session when logging in through SSH:

if [[ -z $TMUX && -n $SSH_TTY ]]; then
  me=$(whoami)

  if tmux has-session -t $me 2>/dev/null; then
    tmux -2 attach-session -t $me
  else
    tmux -2 new-session -s $me
  fi
fi

This seems to work. However, if I exit my tmux session, I am sent back into zsh.

Can I have it so that exit in the tmux session will also end the SSH session? Even better: Could exit in tmux just detach the tmux session, then end the SSH session?

I'm thinking that maybe a zsh alias, defined when [[ -n $TMUX ]], would do the trick, but I'm not sure what would work...

Solution

You can use exec when running tmux. It will give full control to the job you are starting, once that process exits (be it tmux, echo or ls) the shell will exit.

From man zshbuiltins exec ... Replace the current shell with an external command rather than forking.

if [[ -z $TMUX && -n $SSH_TTY ]]; then
  me=$(whoami)

  if tmux has-session -t $me 2>/dev/null; then
     exec tmux -2 attach-session -t $me
  else
     exec tmux -2 new-session -s $me
  fi
fi

the other alternative is to place a shell-script that starts or joins tmux, and ask ssh to run that instead of your shell.

 `alias tmux-ssh="ssh user@target-host -t /home/foo/my-tmux-script`

(the script must obviously be located at your remote host)