worktree script for tmux

worktree.sh

#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<'EOF'
usage:
  worktree [worktree-name]
  worktree --headless worktree-name
  worktree --undo [branch-name]
  worktree --headless --undo branch-name
  worktree -h | --help

Without argument:
  Interactively select an existing worktree and open a tmux window.

With argument:
  Create a new worktree with the given name, check out the branch,
  run setup, and open a new tmux window with Kimi, terminal, and Junie panes.

--headless:
  Create or undo a worktree without any tmux interaction.  Use with a
  worktree name to create only, or with --undo to merge/clean up only.

--undo:
  Merge the worktree branch, close the tmux window, remove the worktree,
  delete the directory, and delete the branch.
EOF
}

die() {
  echo "error: $*" >&2
  exit 1
}

require_cmd() {
  command -v "$1" >/dev/null 2>&1 || die "missing required command: $1"
}

repo_root() {
  # When installed globally, find the repo from the current working directory.
  # Fall back to the script's directory for local/repo-relative usage.
  local root
  root="$(git rev-parse --show-toplevel 2>/dev/null)" || true
  if [[ -n "$root" ]]; then
    echo "$root"
    return 0
  fi
  cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1
  pwd
}

# Run the project's .worktree setup script.
# If .worktree does not exist, create it as an executable script with
# example content so the user can edit it for next time.
run_worktree_script() {
  local root="$1"
  local dest="$2"
  local name="${3:-}"
  local worktree_file="$root/.worktree"

  if [[ ! -f "$worktree_file" ]]; then
    cat > "$worktree_file" <<'INNEREOF'
#!/usr/bin/env bash
# .worktree — setup script run AUTOMATICALLY after creating a new worktree.
#
# This script can do WHATEVER YOU WANT — it is just a shell script sourced with
# the variables below already exported.  Edit it freely; it runs on every
# worktree creation (but is NOT run when opening an existing worktree).
#
# ── exported environment variables ───────────────────────────────────────
#   WORKTREE_NAME   name of the worktree (same as the branch name)
#   TARGET_DIR      absolute path to the new worktree directory
#   REPO_ROOT       absolute path to the main repository root
#   CWD             the directory you were in when you ran "worktree"
# ──────────────────────────────────────────────────────────────────────────
#
# ── examples ──────────────────────────────────────────────────────────────
#
#   # Copy a local secrets file from the repo into the new worktree:
#   cp "$REPO_ROOT/.env.local" "$TARGET_DIR/"
#
#   # Symlink a large shared directory to avoid duplicating disk space:
#   ln -s /mnt/bigdata/models "$TARGET_DIR/models"
#
#   # Download something directly into the worktree:
#   curl -fsSLo "$TARGET_DIR/data/some-file.tar.gz" \
#     https://example.com/some-file.tar.gz
#
#   # Run a host-side script that configures the worktree:
#   bash /opt/scripts/configure-worktree.sh "$TARGET_DIR" "$WORKTREE_NAME"
#
#   # rsync a shared cache from a remote machine:
#   rsync -avz build-cache.local:/cache/ "$TARGET_DIR/.cache/"
#
# ──────────────────────────────────────────────────────────────────────────

# ---- your custom setup starts below ----


# ---- your custom setup ends here ----
INNEREOF
    chmod +x "$worktree_file"
    echo "Created $worktree_file (executable).  Edit it to customise worktree setup." >&2
    return 0
  fi

  if [[ ! -x "$worktree_file" ]]; then
    chmod +x "$worktree_file"
  fi

  WORKTREE_NAME="$name" \
    TARGET_DIR="$dest" \
    REPO_ROOT="$root" \
    CWD="${OLDPWD:-$PWD}" \
    "$worktree_file"
}

list_worktrees() {
  git worktree list --porcelain
}

choose_worktree_dir() {
  local -a wt_dirs=()
  local -a wt_branches=()

  local tty="/dev/tty"
  local has_tty="false"
  if [[ -r "$tty" && -w "$tty" ]]; then
    has_tty="true"
  fi

  local line
  local wt_dir=""
  local wt_branch=""
  while IFS= read -r line; do
    case "$line" in
      worktree\ *)
        wt_dir="${line#worktree }"
        wt_branch=""
        ;;
      branch\ *)
        wt_branch="${line#branch refs/heads/}"
        ;;
      detached)
        wt_branch="(detached)"
        ;;
      "")
        if [[ -n "$wt_dir" ]]; then
          wt_dirs+=("$wt_dir")
          wt_branches+=("${wt_branch:-}" )
          wt_dir=""
          wt_branch=""
        fi
        ;;
    esac
  done < <(list_worktrees)

  # `git worktree list --porcelain` doesn't guarantee a trailing blank line.
  if [[ -n "$wt_dir" ]]; then
    wt_dirs+=("$wt_dir")
    wt_branches+=("${wt_branch:-}")
  fi

  if (( ${#wt_dirs[@]} == 0 )); then
    die "no worktrees found"
  fi

  local i
  for i in "${!wt_dirs[@]}"; do
    local n=$((i + 1))
    local d="${wt_dirs[$i]}"
    local b="${wt_branches[$i]}"
    if [[ -n "$b" ]]; then
      if [[ "$has_tty" == "true" ]]; then
        printf '%2d) %s  [%s]\n' "$n" "$d" "$b" >"$tty"
      else
        printf '%2d) %s  [%s]\n' "$n" "$d" "$b" >&2
      fi
    else
      if [[ "$has_tty" == "true" ]]; then
        printf '%2d) %s\n' "$n" "$d" >"$tty"
      else
        printf '%2d) %s\n' "$n" "$d" >&2
      fi
    fi
  done

  local choice
  if [[ "$has_tty" == "true" ]]; then
    printf 'choose worktree (1-%d): ' "${#wt_dirs[@]}" >"$tty"
    IFS= read -r choice <"$tty"
  else
    printf 'choose worktree (1-%d): ' "${#wt_dirs[@]}" >&2
    IFS= read -r choice
  fi
  [[ "$choice" =~ ^[0-9]+$ ]] || die "not a number: $choice"
  (( choice >= 1 && choice <= ${#wt_dirs[@]} )) || die "out of range: $choice"

  echo "${wt_dirs[$((choice - 1))]}"
}

choose_worktree_branch_interactive() {
  local -a wt_dirs=()
  local -a wt_branches=()

  local tty="/dev/tty"
  local has_tty="false"
  if [[ -r "$tty" && -w "$tty" ]]; then
    has_tty="true"
  fi

  local line
  local wt_dir=""
  local wt_branch=""
  while IFS= read -r line; do
    case "$line" in
      worktree\ *)
        wt_dir="${line#worktree }"
        wt_branch=""
        ;;
      branch\ *)
        wt_branch="${line#branch refs/heads/}"
        ;;
      detached)
        wt_branch="(detached)"
        ;;
      "")
        if [[ -n "$wt_dir" && -n "$wt_branch" && "$wt_branch" != "(detached)" ]]; then
          wt_dirs+=("$wt_dir")
          wt_branches+=("$wt_branch")
        fi
        wt_dir=""
        wt_branch=""
        ;;
    esac
  done < <(list_worktrees)

  # Handle last entry if file doesn't end with blank line
  if [[ -n "$wt_dir" && -n "$wt_branch" && "$wt_branch" != "(detached)" ]]; then
    wt_dirs+=("$wt_dir")
    wt_branches+=("$wt_branch")
  fi

  if (( ${#wt_dirs[@]} == 0 )); then
    die "no worktrees with branches found"
  fi

  local i
  for i in "${!wt_dirs[@]}"; do
    local n=$((i + 1))
    local d="${wt_dirs[$i]}"
    local b="${wt_branches[$i]}"
    if [[ "$has_tty" == "true" ]]; then
      printf '%2d) %s  [%s]\n' "$n" "$d" "$b" >"$tty"
    else
      printf '%2d) %s  [%s]\n' "$n" "$d" "$b" >&2
    fi
  done

  local choice
  if [[ "$has_tty" == "true" ]]; then
    printf 'choose worktree to unwork (1-%d): ' "${#wt_dirs[@]}" >"$tty"
    IFS= read -r choice <"$tty"
  else
    printf 'choose worktree to unwork (1-%d): ' "${#wt_dirs[@]}" >&2
    IFS= read -r choice
  fi
  [[ "$choice" =~ ^[0-9]+$ ]] || die "not a number: $choice"
  (( choice >= 1 && choice <= ${#wt_dirs[@]} )) || die "out of range: $choice"

  echo "${wt_branches[$((choice - 1))]}"
}

cmd_work() {
  require_cmd git

  local headless="${1:-false}"
  shift || true

  if [[ "$headless" != "true" ]]; then
    require_cmd tmux
  fi

  # Make sure git commands run from within the repository root.
  local root
  root="$(repo_root)"
  cd "$root"

  local worktree_name="${1:-}"
  local chosen_dir
  local branch_name

  if [[ -n "$worktree_name" ]]; then
    # Create new worktree in parent directory
    local parent_dir
    parent_dir="$(dirname "$root")"
    chosen_dir="${parent_dir}/${worktree_name}"

    # Check if worktree already exists
    if [[ -d "$chosen_dir" ]]; then
      die "worktree already exists: $chosen_dir"
    fi

    # Check if branch exists, create if not
    if ! git show-ref --verify --quiet "refs/heads/${worktree_name}"; then
      echo "Creating new branch: ${worktree_name}"
      git branch "${worktree_name}"
    fi

    echo "Creating worktree at: ${chosen_dir}"
    git worktree add "$chosen_dir" "$worktree_name"

    # Run the project's .worktree setup script
    run_worktree_script "$root" "$chosen_dir" "$worktree_name"

    branch_name="$worktree_name"
  else
    # Interactive mode: choose existing worktree
    chosen_dir="$(choose_worktree_dir)"
    [[ -d "$chosen_dir" ]] || die "worktree dir does not exist: $chosen_dir"

    branch_name="$(git -C "$chosen_dir" rev-parse --abbrev-ref HEAD)"
    worktree_name="$branch_name"
  fi

  # If headless, stop here — no tmux interaction
  if [[ "$headless" == "true" ]]; then
    echo "Headless worktree created: $chosen_dir"
    return 0
  fi

  # Verify running inside a tmux session
  [[ -z "${TMUX:-}" ]] && die "not running inside a tmux session"

  local session
  session="$(tmux display-message -p '#S')"

  # If a window with this name already exists, switch to it and exit
  if tmux list-windows -t "$session" -F '#{window_name}' | grep -qx "$worktree_name"; then
    tmux select-window -t "$session:$worktree_name"
    return 0
  fi

  # Create a new window in the current session with Kimi
  tmux new-window -n "$worktree_name" -t "$session" \
    "bash -lc 'cd \"$chosen_dir\" && runkimi --yolo || true; exec bash -l'"

  # Split pane 0 vertically (bottom 30%) for a terminal pane
  tmux split-window -v -p 30 -t "$session:$worktree_name.0" \
    "bash -lc 'cd \"$chosen_dir\" && exec bash -l'"

  # Split pane 0 horizontally (top right) for the junie pane
  tmux split-window -h -t "$session:$worktree_name.0" \
    "bash -lc 'cd \"$chosen_dir\" && junie --session-id=\"$worktree_name\" --brave || true; exec bash -l'"

  # Select the top-left pane (kimi)
  tmux select-pane -t "$session:$worktree_name.0"
}

cmd_undo() {
  require_cmd git

  local headless="${1:-false}"
  shift || true

  if [[ "$headless" != "true" ]]; then
    require_cmd tmux
  fi

  # Make sure git commands run from within the repository root.
  local root
  root="$(repo_root)"
  cd "$root"

  local branch_name="${1:-}"

  if [[ -z "$branch_name" ]]; then
    # Interactive mode
    branch_name="$(choose_worktree_branch_interactive)"
  fi

  # Find the worktree directory for this branch
  local worktree_dir=""
  local line
  local wt_dir=""
  local wt_branch=""
  while IFS= read -r line; do
    case "$line" in
      worktree\ *)
        wt_dir="${line#worktree }"
        wt_branch=""
        ;;
      branch\ *)
        wt_branch="${line#branch refs/heads/}"
        ;;
      "")
        if [[ "$wt_branch" == "$branch_name" ]]; then
          worktree_dir="$wt_dir"
          break
        fi
        wt_dir=""
        wt_branch=""
        ;;
    esac
  done < <(list_worktrees)

  # Handle last entry
  if [[ -z "$worktree_dir" && "$wt_branch" == "$branch_name" ]]; then
    worktree_dir="$wt_dir"
  fi

  if [[ -z "$worktree_dir" ]]; then
    die "no worktree found for branch: $branch_name"
  fi

  # Prevent operating on the main worktree (repo root)
  if [[ "$worktree_dir" == "$root" ]]; then
    die "refusing to operate on the main repository worktree"
  fi

  # Merge the branch; if merge fails, exit before touching the tmux window
  git merge "$branch_name" || die "merge failed for branch: $branch_name"

  # Close the tmux window gracefully (skip if headless)
  if [[ "$headless" != "true" ]]; then
    local session
    session="$(tmux display-message -p '#S')"

    if tmux list-windows -t "$session" -F '#{window_name}' | grep -qx "$branch_name"; then
      # Send /quit to pane 0 (kimi) and pane 2 (junie)
      tmux send-keys -t "$session:$branch_name.0" "/quit" Enter
      tmux send-keys -t "$session:$branch_name.2" "/quit" Enter

      # Give agents a moment to start shutting down
      sleep 1

      # Kill the window
      tmux kill-window -t "$session:$branch_name"
    else
      echo "tmux window '$branch_name' not found, continuing"
    fi
  fi

  echo "Removing worktree: $worktree_dir"
  git worktree remove --force "$worktree_dir"

  echo "Deleting directory: $worktree_dir"
  rm -rf "$worktree_dir"

  echo "Deleting branch: $branch_name"
  git branch -d "$branch_name"

  echo "Successfully undone: $branch_name"
}

main() {
  local headless="false"

  # Parse --headless flag (must come before subcommand/name)
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --headless)
        headless="true"
        shift
        ;;
      --undo)
        shift
        cmd_undo "$headless" "$@"
        return
        ;;
      -h|--help|help)
        usage
        ;;
      -*)
        usage
        exit 2
        ;;
      *)
        cmd_work "$headless" "$@"
        return
        ;;
    esac
  done

  # No arguments at all — interactive mode
  cmd_work "$headless"
}

main "$@"