Rework code shim

home/.chezmoiscripts/run_after_90-install-vscode-extensions.sh.tmpl

@@ -8,7 +8,7 @@ true || source ../.chezmoitemplates/scripts-library
 true || source ../.chezmoitemplates/vscode-library
 
 function code() {
-  USE_NATIVE_CODE=1 command code "$@"
+  NO_REMOTE=1 command code "$@"
 }
 
 install_vscode_extensions

home/dot_local/bin/executable_code

@@ -2,182 +2,241 @@
 
 set -eu
 
-get_in_path_except_current() {
-  entries=$(which -a "$1") &&
-    echo "${entries}" | grep -A1 "$0" | grep -v "$0" | head -1 | grep .
-}
+if [ -n "${DEBUG:-}" ]; then
+  set -x
+fi
 
-is_path_from_windows() {
-  echo "${1}" | grep -q '^/mnt/.*$'
-}
+# Find whether we are trying to open a directory
+target_dir=""
+for arg in "$@"; do
+  # Skip options like --wait
+  if [ "${arg#-}" = "${arg}" ]; then
+    if [ -d "${arg}" ]; then
+      target_dir="${arg}"
+    fi
+    break
+  fi
+done
+unset arg
 
-get_in_path_from_windows() {
-  entries=$(which -a "$1") &&
-    echo "${entries}" | grep -v -A99 "$0" | grep -v "$0" | grep -m1 '^/mnt/.*$' | grep .
-}
+# If we are trying to open a folder, drop the --wait flag to work around:
+# https://github.com/twpayne/chezmoi/issues/1068
+for arg in "$@"; do
+  shift
 
-is_first_arg_a_folder() {
-  for arg; do
-    shift
-    # Skip options like --wait
-    if [ "${arg#-}" = "${arg}" ]; then
-      if [ -d "${arg}" ]; then
-        echo "${arg}"
-        return 0
-      else
-        return 1
-      fi
+  if [ "${arg}" = "--wait" ] || [ "${arg}" = "-w" ]; then
+    if [ -n "${target_dir}" ]; then
+      echo "INFO(dotfiles-code): dropping --wait since a directory is trying to be opened" >&2
+      continue
     fi
-  done
-  return 1
-}
+  fi
 
-if [ -n "${DEBUG:-}" ]; then
-  set -x
+  set -- "$@" "${arg}"
+done
+unset arg
+
+# Decide whether we want to try opening remote sessions or not
+remote=true
+if [ -n "${NO_REMOTE:-}" ]; then
+  remote=false
 fi
 
-# Used to work around the issue: https://github.com/twpayne/chezmoi/issues/1068
-if is_first_arg_a_folder "$@" >/dev/null; then
-  # Removes -w and --wait flags
-  for arg; do
-    shift
-    if [ "${arg}" = "--wait" ] || [ "${arg}" = "-w" ]; then
-      echo "dropping --wait or -w flag" >&2
-      continue
-    fi
-    set -- "$@" "${arg}"
-  done
+is_wsl=false
+if [ -n "${WSL_DISTRO_NAME:-}" ] || [ -n "${IS_WSL:-}" ]; then
+  is_wsl=true
 fi
 
-# Check for a --devcontainer flag
-for arg; do
+# Handle --devcontainer flag
+devcontainer=false
+for arg in "$@"; do
   shift
+
   if [ "${arg}" = "--devcontainer" ]; then
-    if {
-      workspace_folder=$(is_first_arg_a_folder "$@") &&
-        workspace_folder="$(realpath "${workspace_folder}")" &&
-        [ -d "${workspace_folder}" ]
-    }; then
-      if ! workspace_folder_in_devcontainer=$(
+    devcontainer=true
+
+    if [ -n "${target_dir}" ]; then
+      target_dir=$(realpath "${target_dir}")
+
+      if ! devcontainer_workspace_folder=$(
         # https://github.com/microsoft/vscode-remote-release/issues/2133#issuecomment-1430651840
-        devcontainer read-configuration --workspace-folder "${workspace_folder}" 2>/dev/null |
-          jq --exit-status --raw-output .workspace.workspaceFolder |
-          grep .
+        devcontainer read-configuration --workspace-folder "${target_dir}" 2>/dev/null |
+          jq --exit-status --raw-output .workspace.workspaceFolder
       ); then
-        echo "failed to read the devcontainer workspace folder." >&2
+        echo "ERROR(dotfiles-code): failed to read the devcontainer workspace folder" >&2
         exit 1
       fi
 
-      if [ -z "${USE_NATIVE_CODE:-}" ] && [ -n "${WSL_DISTRO_NAME:-}" ]; then
-        workspace_folder=$(wslpath -w "${workspace_folder}")
+      if [ "${remote}" = true ] && [ "${is_wsl}" = true ]; then
+        target_dir=$(wslpath -w "${target_dir}")
       fi
-      path_id=$(printf "%s" "${workspace_folder}" | xxd -ps -c 256)
 
-      set -- --folder-uri "vscode-remote://dev-container%2B${path_id}${workspace_folder_in_devcontainer}"
+      devcontainer_path_id=$(printf "%s" "${target_dir}" | xxd -ps -c 256)
       break
     else
-      echo "when using --devcontainer, the first arg must be a folder."
+      echo "ERROR(dotfiles-code): when using --devcontainer, the first argument must be a directory"
       exit 1
     fi
   fi
+
   set -- "$@" "${arg}"
 done
+unset arg
+
+exec_code() {
+  code="$1"
+  shift
+
+  if [ "${devcontainer}" = false ]; then
+    exec "${code}" "$@"
+  fi
+
+  devcontainer_remote_id=""
+  if [ -n "${ssh_host:-}" ]; then
+    devcontainer_remote_id="@ssh-remote%2B${ssh_host}"
+  elif [ "${is_wsl}" = false ] && [ -n "${VSCODE_IPC_HOOK_CLI:-}" ]; then
+    # Try to find remote tunnel host
+    if tunnel_host=$(
+      "${code}" --status 2>/dev/null | grep -oP "^Remote:[ ]+\K.+" | grep -v '[ :]' | head -1 | grep .
+    ); then
+      devcontainer_remote_id="@tunnel%2B${tunnel_host}"
+    else
+      echo "ERROR(dotfiles-code): failed to get the remote id" >&2
+      exit 1
+    fi
+  fi
+  exec "${code}" --folder-uri "vscode-remote://dev-container%2B${devcontainer_path_id}${devcontainer_remote_id}${devcontainer_workspace_folder}"
+}
 
-# Check if we are running in SSH, unless USE_NATIVE_CODE is set
-if [ -z "${USE_NATIVE_CODE:-}" ] && who -m | grep -q .; then
-  # We will try to find an existing connection and reuse it.
-  # This allows opening VS Code windows from a SSH remote
-  # session even out of the VS Code integrated terminal.
-  printf "searching for remote ssh session (set USE_NATIVE_CODE=1 to use code from Linux)..." >&2
-  no_sockets_message=" no sockets found."
-  if [ -d "${HOME}/.vscode-server/bin" ]; then
-
-    # Find most recent folder inside of ~/.vscode-server/bin
-    if bin_folder=$(find "${HOME}/.vscode-server/bin" \
-      -mindepth 1 -maxdepth 1 -type d -printf "%T@ %p\n" |
-      sort -n | cut -d' ' -f 2- | tail -n 1 | grep .); then
-
-      # Ensure code binary is executable
-      bin_path="${bin_folder}/bin/remote-cli/code"
-      if [ -x "${bin_path}" ]; then
-
-        # Find most recent sock
+vscode_server_dir="${HOME}/.vscode-server"
+
+# Handle vscode ssh, unless NO_REMOTE is set
+if [ "${remote}" = true ] && [ -z "${VSCODE_IPC_HOOK_CLI:-}" ] && who -m | grep -q .; then
+  # We will try to find an existing connection and reuse it. This allows
+  # opening VS Code windows from a SSH remote session even out of the VS Code
+  # integrated terminal.
+  # PS: this only works if there is at least one VS Code window already
+  # connected to this machine.
+  echo "INFO(dotfiles-code): ssh session detected, searching for vscode remote sessions (set NO_REMOTE=1 to disable this)" >&2
+  no_sockets_message="no vscode remote sessions found"
+
+  code_search_dir="${vscode_server_dir}/cli/servers"
+  if [ -d "${code_search_dir}" ]; then
+    # Find the most recent directory inside of code_search_dir
+    if code_search_dir=$(
+      find "${code_search_dir}" \
+        -mindepth 1 -maxdepth 1 -type d -printf "%T@ %p\n" |
+        sort -nr | cut -d' ' -f 2- | head -1 | grep .
+    ); then
+      # Ensure the bin is valid
+      code="${code_search_dir}/server/bin/remote-cli/code"
+      if [ -x "${code}" ]; then
+        # Find all vscode-ipc-*.sock files, trying to connect to the newest first
         uid=$(id -u)
-        if sock_paths=$(find "/run/user/${uid}/" \
-          -mindepth 1 -maxdepth 1 -type s -name "vscode-ipc-*.sock" -printf "%T@ %p\n" |
-          sort --numeric-sort --reverse | cut -d' ' -f 2- | grep .); then
-
-          sockets_count=$(echo "${sock_paths}" | wc -l)
-          echo " ${sockets_count} sockets found." >&2
-          no_sockets_message=" no more sockets found."
-          echo "trying to connect..." >&2
-          for sock_path in ${sock_paths}; do
-            if VSCODE_IPC_HOOK_CLI="${sock_path}" "${bin_path}" "$@"; then
-              exit 0
+        if sockets=$(
+          find "/run/user/${uid}/" \
+            -mindepth 1 -maxdepth 1 -type s -name "vscode-ipc-*.sock" -printf "%T@ %p\n" |
+            sort -nr | cut -d' ' -f 2- | grep .
+        ); then
+          sockets_count=$(echo "${sockets}" | wc -l)
+          echo "INFO(dotfiles-code): ${sockets_count} vscode remote sessions found" >&2
+          unset sockets_count
+
+          for socket in ${sockets}; do
+            echo "INFO(dotfiles-code): trying to connect to ${socket}" >&2
+            export VSCODE_IPC_HOOK_CLI="${socket}"
+            if ssh_host=$(
+              timeout 5s "${code}" --status 2>/dev/null | grep -m1 -oP "^Remote:[ ]+SSH:[ ]+\K.+"
+            ); then
+              echo "INFO(dotfiles-code): worked, using it" >&2
+              exec_code "${code}" "$@"
             else
-              echo "removing socket and trying next..." >&2
-              rm -f "${sock_path}"
+              echo "INFO(dotfiles-code): did not work" >&2
             fi
           done
+
+          no_sockets_message="no more vscode remote sessions found"
+          unset socket VSCODE_IPC_HOOK_CLI output
         fi
+        unset uid sockets
       fi
+      unset code
     fi
   fi
-  echo "${no_sockets_message}" >&2
-else
-  # shellcheck disable=SC2310
-  if code="$(get_in_path_except_current code)"; then
-    # shellcheck disable=SC2310
-    if { [ -n "${WSL_DISTRO_NAME:-}" ] || [ -n "${IS_WSL:-}" ]; } &&
-      ! is_path_from_windows "${code}"; then
-      echo >&2
-      if get_in_path_from_windows code >/dev/null && [ -z "${USE_NATIVE_CODE:-}" ]; then
-        echo "using code from Windows (set USE_NATIVE_CODE=1 to use code from Linux)" >&2
-        code="$(get_in_path_from_windows code)"
-      else
-        echo "using code from Linux (unset USE_NATIVE_CODE to use code from Windows)" >&2
-        export DONT_PROMPT_WSL_INSTALL=1
-      fi
-    fi
+  echo "WARNING(dotfiles-code): ${no_sockets_message}" >&2
+  unset code_search_dir
+fi
 
-    exec "${code}" "$@"
+# Only consider next "code"s in PATH
+this_script=$(realpath "$0")
+executables=$(
+  # shellcheck disable=SC2230
+  which -a code | grep -A99 "${this_script}" | grep -v "${this_script}" || true
+)
+
+if [ "${remote}" = false ]; then
+  echo "INFO(dotfiles-code): using code from linux" >&2
+  # Avoids reopening the remote session if called from within the remote session
+  unset VSCODE_IPC_HOOK_CLI
+  if [ "${is_wsl}" = true ]; then
+    export DONT_PROMPT_WSL_INSTALL=1
+    executables=$(echo "${executables}" | grep -v '^/mnt/.*$' || true)
   fi
+  executables=$(echo "${executables}" | grep -v "${vscode_server_dir}" || true)
+elif [ "${is_wsl}" = true ]; then
+  echo "INFO(dotfiles-code): using code from windows (set NO_REMOTE=1 to disable this)" >&2
+  executables=$(echo "${executables}" | grep '^/mnt/.*$' || true)
+fi
 
-  echo >&2
-  printf "code is not available, " >&2
-  if command -v code-insiders >/dev/null 2>&1; then
-    echo "using code-insiders instead" >&2
-    exec code-insiders "$@"
-  fi
+if code=$(echo "${executables}" | head -1 | grep .); then
+  exec_code "${code}" "$@"
 fi
 
-if is_first_arg_a_folder "$@" >/dev/null; then
-  echo "not falling back to nano, vim or vi since you are trying to open a folder." >&2
+# When NO_REMOTE is set, we should not try to fallback because the user has
+# explicitly requested an intention to use code.
+if [ "${remote}" = false ]; then
+  echo "ERROR(dotfiles-code): code is not available" >&2
   exit 127
 fi
 
-if command -v nano >/dev/null 2>&1; then
+echo "WARNING(dotfiles-code): code is not available" >&2
+if command -v code-insiders >/dev/null; then
+  echo "INFO(dotfiles-code): using code-insiders instead" >&2
+  exec_code code-insiders "$@"
+fi
+
+if [ -n "${target_dir}" ]; then
+  echo "ERROR(dotfiles-code): not falling back to nano, vim, or vi since a directory is trying to be opened" >&2
+  exit 127
+fi
+
+if command -v nano >/dev/null; then
   editor="nano"
-elif command -v vim >/dev/null 2>&1; then
+elif command -v vim >/dev/null; then
   editor="vim"
-elif command -v vi >/dev/null 2>&1; then
+elif command -v vi >/dev/null; then
   editor="vi"
 else
-  echo "and neither code-insiders, nano, vim or vi." >&2
+  echo "ERROR(dotfiles-code): neither code-insiders, nano, vim, or vi is available" >&2
   exit 127
 fi
 
-echo "using ${editor} instead" >&2
+echo "INFO(dotfiles-code): using ${editor} instead" >&2
 
-# Removes -w and --wait flags
-# code duplication here is required because there is no way to set
-# script level args inside of functions in POSIX (on Bash we can use arrays)
-for arg; do
+for arg in "$@"; do
   shift
+
+  # Remove the --wait flag as it's not supported by nano, vim, or vi
   if [ "${arg}" = "--wait" ] || [ "${arg}" = "-w" ]; then
-    echo "dropping --wait or -w flag" >&2
+    echo "INFO(dotfiles-code): dropping --wait as it's not supported by ${editor}" >&2
     continue
   fi
+
+  # If any other option is passed, we fail
+  if [ "${arg#-}" != "${arg}" ]; then
+    echo "ERROR(dotfiles-code): not falling back to ${editor} since ${arg} was passed" >&2
+    exit 127
+  fi
+
   set -- "$@" "${arg}"
 done