feat: add windmill-ee-private worktree support to workmux (#8034)

* feat: add windmill-ee-private worktree support to workmux

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* feat: add EE worktree cleanup on remove and parent-dir lookup

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

.workmux.yaml

@@ -26,6 +26,9 @@ auto_name:
 post_create:
   - ./scripts/worktree-env
 
+pre_remove:
+  - ./scripts/worktree-cleanup
+
 panes:
   - command: <agent>
     focus: true

backend/substitute_ee_code.sh

@@ -7,6 +7,7 @@ REVERT="NO"
 COPY="NO"
 MOVE_NEW_FILES="NO"
 EE_CODE_DIR="../windmill-ee-private/"
+DIR_EXPLICIT="NO"
 
 while [[ $# -gt 0 ]]; do
   case $1 in
@@ -34,6 +35,7 @@ while [[ $# -gt 0 ]]; do
       # Path to the local directory of the windmill-ee-private repository. By defaults, it
       # assumes it is cloned next to the Windmill OSS repo.
       EE_CODE_DIR="$2"
+      DIR_EXPLICIT="YES"
       shift # past argument
       shift # past value
       ;;
@@ -55,10 +57,22 @@ else
 fi
 
 # Fallback to ~/windmill-ee-private if the default location doesn't exist
-if [ ! -d "${EE_CODE_DIR}" ] && [ "${EE_CODE_DIR}" == "${root_dirpath}/../windmill-ee-private/" ]; then
+if [ ! -d "${EE_CODE_DIR}" ]; then
   EE_CODE_DIR="${HOME}/windmill-ee-private"
 fi
 
+# Unless --dir was explicitly set, try to find an EE worktree on the same branch
+if [ "$DIR_EXPLICIT" == "NO" ] && [ -d "${HOME}/windmill-ee-private" ]; then
+  current_branch=$(git -C "${root_dirpath}" branch --show-current 2>/dev/null || true)
+  if [ -n "$current_branch" ]; then
+    ee_worktree=$(git -C "${HOME}/windmill-ee-private" worktree list 2>/dev/null \
+      | awk -v branch="[${current_branch}]" '$NF == branch {print $1; exit}')
+    if [ -n "$ee_worktree" ] && [ -d "$ee_worktree" ]; then
+      EE_CODE_DIR="$ee_worktree"
+    fi
+  fi
+fi
+
 echo "EE code directory = ${EE_CODE_DIR} | Revert = ${REVERT}"
 
 if [ ! -d "${EE_CODE_DIR}" ]; then

scripts/worktree-cleanup

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Remove the matching windmill-ee-private worktree if one exists
+wt_basename=$(basename "$(pwd)")
+
+# Check parent directory first (sibling to worktree root), then fall back to home
+parent_dir="$(cd "$(pwd)/.." && pwd)"
+if [ -d "${parent_dir}/windmill-ee-private" ]; then
+  ee_repo="${parent_dir}/windmill-ee-private"
+else
+  ee_repo="${HOME}/windmill-ee-private"
+fi
+
+ee_worktree_dir="${ee_repo}__worktrees/${wt_basename}"
+if [ -d "$ee_worktree_dir" ]; then
+  git -C "$ee_repo" worktree remove "$ee_worktree_dir" --force 2>/dev/null \
+    && echo "Removed EE worktree at $ee_worktree_dir" \
+    || echo "Warning: Could not remove EE worktree at $ee_worktree_dir"
+fi

scripts/worktree-env

@@ -49,3 +49,35 @@ REMOTE=http://localhost:$backend_port
 EOF
 
 echo "Created .env.local with ports: backend=$backend_port, frontend=$frontend_port"
+
+# --- Create matching windmill-ee-private worktree ---
+# Check parent directory first (sibling to worktree root), then fall back to home
+parent_dir="$(cd "$(pwd)/.." && pwd)"
+if [ -d "${parent_dir}/windmill-ee-private" ]; then
+  ee_repo="${parent_dir}/windmill-ee-private"
+else
+  ee_repo="${HOME}/windmill-ee-private"
+fi
+if [ -d "$ee_repo" ]; then
+  branch=$(git branch --show-current 2>/dev/null || true)
+  wt_basename=$(basename "$(pwd)")
+  ee_worktree_dir="${ee_repo}__worktrees/${wt_basename}"
+
+  if [ -n "$branch" ] && [ ! -d "$ee_worktree_dir" ]; then
+    mkdir -p "$(dirname "$ee_worktree_dir")"
+
+    # Fetch latest so we can check out remote branches
+    git -C "$ee_repo" fetch --quiet 2>/dev/null || true
+
+    # Try: existing branch, then new branch from main
+    if git -C "$ee_repo" worktree add "$ee_worktree_dir" "$branch" 2>/dev/null; then
+      echo "Created EE worktree at $ee_worktree_dir (branch: $branch)"
+    elif git -C "$ee_repo" worktree add -b "$branch" "$ee_worktree_dir" main 2>/dev/null; then
+      echo "Created EE worktree at $ee_worktree_dir (new branch: $branch from main)"
+    else
+      echo "Warning: Could not create EE worktree for branch $branch"
+    fi
+  elif [ -d "$ee_worktree_dir" ]; then
+    echo "EE worktree already exists at $ee_worktree_dir"
+  fi
+fi