Put this in your .bash_profile(s) to open VS Code on your local machine when SSH into a remote machine.

.bash_profile

# this needs to be on the *remote* machine.

function code() {
  if [ -z "$SSH_CLIENT" ]
  then
    unset -f code
    code $@
    return
  fi
  
  ssh_client_host=$(echo $SSH_CLIENT | cut -d ' ' -f1)
  
  if [ -z "$1" ]
  then
    argpath="."
  else
    argpath="$1"
  fi
  
  # ssh-remote requires a directory
  resolved=$(realpath -m $argpath)
  if [ -f "$resolved" ]; then
    resolved=$(dirname $resolved)
  fi
  ssh $ssh_client_host code \
    --folder-uri=vscode-remote://ssh-remote+$(hostname)$resolved
  unset resolved
  unset argpath
  unset ssh_client_host
}