Automatically update $DISPLAY for each tmux pane after attaching to session

tmux.conf

set-hook -g client-attached 'run-shell /bin/update_display.sh'

update_display.sh

#!/usr/bin/env bash

# The problem:
# When you `ssh -X` into a machine and attach to an existing tmux session, the session
# contains the old $DISPLAY env variable. In order the x-server/client to work properly,
# you have to update $DISPLAY after connection. For example, the old $DISPLAY=:0 and
# you need to change to DISPLAY=localhost:10.0 for the ssh session to
# perform x-forwarding properly.

# The solution:
# When attaching to tmux session, update $DISPLAY for each tmux pane in that session
# This is performed by using tmux send-keys to the shell.
# This script handles updating $DISPLAY within vim also

NEW_DISPLAY=$DISPLAY
# NEW_DISPLAY=$(tmux show-env | sed -n 's/^DISPLAY=//p')

# Update $DISPLAY in bash, zsh and vim/nvim
tmux list-panes -s -F "#{session_name}:#{window_index}.#{pane_index} #{pane_current_command}" | \
while read pane_process
do
    IFS=' ' read -ra pane_process <<< "$pane_process"
    if [[ "${pane_process[1]}" == "zsh" || "${pane_process[1]}" == "bash" ]]; then
        tmux send-keys -t ${pane_process[0]} "export DISPLAY=$NEW_DISPLAY" Enter
    elif [[ "${pane_process[1]}" == *"vi"* ]]; then
        tmux send-keys -t ${pane_process[0]} Escape
        tmux send-keys -t ${pane_process[0]} ":let \$DISPLAY = \"$NEW_DISPLAY\"" Enter
        tmux send-keys -t ${pane_process[0]} ":silent! xrestore" Enter
    fi
done

@qjzh0603, glad you got it working.
You're right, I experience the same annoyance. There isn't a good automated solution to this problem that I know of, so I just created an alias to update my display manually after I quit vim:
alias ud='export DISPLAY="`tmux show-env | sed -n \"s/^DISPLAY=//p\"`"'

Thanks for the script!. BTW: I added hook to session switching:
set-hook -g client-session-changed 'run-shell /bin/update_display.sh'
It works. However, DISPLAY is updated every time I switch session event if it's not necessary. Tried to guard the update script with
if [[ $NEW_DISPLAY != $DISPLAY ]]; then
It seems that this condition always satisfies even if DISPLAY really needs update. $DISPLAY gives the value of the shell before attaching to tmux session. So the condition is matched.

Thanks for the script!. BTW: I added hook to session switching: set-hook -g client-session-changed 'run-shell /bin/update_display.sh' It works. However, DISPLAY is updated every time I switch session event if it's not necessary. Tried to guard the update script with if [[ $NEW_DISPLAY != $DISPLAY ]]; then It seems that this condition always satisfies even if DISPLAY really needs update. $DISPLAY gives the value of the shell before attaching to tmux session. So the condition is matched.

@robertbu, I like it! I have updated my script with this logic.

@robertbu, actually that new logic didn't work. When update-display.sh evaluated $DISPLAY, it was already seeing the new DISPLAY value.
In fact, I simplified the script to use: NEW_DISPLAY=$DISPLAY
I'm not sure if you had to do anything else to make this conditional logic work on your side. Would love to get it working too.

Thanks for this. I made a version which updates WAYLAND_DISPLAY based on this: https://github.com/graham33/scripts/blob/master/update_wayland_display.sh. It's just a search and replace on your script, ideally it would be great to have one script that handled both, or maybe any env var in the tmux update-environment list.