Skills › Compound Engineering › ce-optimize › experiment-worktree.sh
experiment-worktree.sh
330 lines · skills/ce-optimize/scripts/experiment-worktree.sh at 3449230
1#!/bin/bash
3# Experiment Worktree Manager
4# Creates, cleans up, and manages worktrees for optimization experiments.
5# Each experiment gets an isolated worktree with copied shared resources.
6#
7# Usage:
8# experiment-worktree.sh create <spec_name> <exp_index> <base_branch> [shared_file ...]
9# experiment-worktree.sh cleanup <spec_name> <exp_index>
10# experiment-worktree.sh cleanup-all <spec_name>
11# experiment-worktree.sh count
12#
13# Worktrees are created at: .worktrees/optimize-<spec>-exp-<NNN>/
14# Branches are named: optimize-exp/<spec>/exp-<NNN>
16set -euo pipefail
18RED='\033[0;31m'
19GREEN='\033[0;32m'
20YELLOW='\033[1;33m'
21BLUE='\033[0;34m'
22NC='\033[0m'
24GIT_ROOT=$(git rev-parse --show-toplevel 2>/dev/null) || {
25 echo -e "${RED}Error: Not in a git repository${NC}" >&2
26 exit 1
27}
29WORKTREE_DIR="$GIT_ROOT/.worktrees"
31# Every path this script deletes is built from these two inputs, so they are
32# checked before any path is composed: a name or shared path that can climb
33# out of its directory would turn a cleanup into a delete of something else.
34require_spec_name() {
35 if [[ ! "$1" =~ ^[a-z0-9][a-z0-9-]*$ ]]; then
36 echo -e "${RED}Error: spec name must be lowercase kebab-case: $1${NC}" >&2
37 exit 1
38 fi
39}
41require_relative_inside() {
42 case "/$1/" in
43 //*|*/../*|*/./*)
44 echo -e "${RED}Error: shared path must be relative and stay inside the repo: $1${NC}" >&2
45 exit 1
46 ;;
47 esac
48}
50# Last-resort removal when git cannot remove a worktree. Only ever a direct
51# child of .worktrees that this script named.
52remove_experiment_dir() {
53 local target="$1"
54 if [[ "$(dirname "$target")" != "$WORKTREE_DIR" ]] || [[ "$(basename "$target")" != optimize-*-exp-* ]] || [[ -L "$target" ]]; then
55 echo -e "${RED}Error: refusing to remove a path outside $WORKTREE_DIR: $target${NC}" >&2
56 return 1
57 fi
58 rm -rf "$target"
59}
61experiment_branch_name() {
62 local spec_name="${1:?Error: spec_name required}"
63 local padded_index="${2:?Error: padded_index required}"
65 # Keep experiment refs outside optimize/<spec> so they do not collide
66 # with the long-lived optimization branch namespace.
67 echo "optimize-exp/${spec_name}/exp-${padded_index}"
68}
70ensure_worktree_exclude() {
71 local exclude_file
72 exclude_file=$(git rev-parse --git-path info/exclude)
74 mkdir -p "$(dirname "$exclude_file")"
76 if ! grep -q "^\.worktrees$" "$exclude_file" 2>/dev/null; then
77 echo ".worktrees" >> "$exclude_file"
78 fi
79}
81is_registered_worktree() {
82 local worktree_path="${1:?Error: worktree_path required}"
84 git worktree list --porcelain | awk -v target="$worktree_path" '
85 $1 == "worktree" && $2 == target { found = 1 }
86 END { exit(found ? 0 : 1) }
87 '
88}
90is_branch_checked_out() {
91 local branch_name="${1:?Error: branch_name required}"
92 local branch_ref="refs/heads/$branch_name"
94 git worktree list --porcelain | awk -v target="$branch_ref" '
95 $1 == "branch" && $2 == target { found = 1 }
96 END { exit(found ? 0 : 1) }
97 '
98}
100reset_worktree_to_base() {
101 local worktree_path="${1:?Error: worktree_path required}"
102 local branch_name="${2:?Error: branch_name required}"
103 local base_branch="${3:?Error: base_branch required}"
104 local current_branch
106 current_branch=$(git -C "$worktree_path" symbolic-ref --quiet --short HEAD 2>/dev/null || true)
107 if [[ "$current_branch" != "$branch_name" ]]; then
108 echo -e "${RED}Error: Existing worktree is on unexpected branch: ${current_branch:-detached} (expected $branch_name)${NC}" >&2
109 echo -e "${RED}Clean up the stale worktree before rerunning this experiment.${NC}" >&2
110 return 1
111 fi
113 echo -e "${YELLOW}Resetting existing experiment worktree to base: $branch_name -> $base_branch${NC}" >&2
114 git -C "$worktree_path" reset --hard "$base_branch" >/dev/null
115 git -C "$worktree_path" clean -fdx >/dev/null
116}
118# Create an experiment worktree
119create_worktree() {
120 local spec_name="${1:?Error: spec_name required}"
121 local exp_index="${2:?Error: exp_index required}"
122 local base_branch="${3:?Error: base_branch required}"
123 shift 3
124 require_spec_name "$spec_name"
125 local shared_file
126 for shared_file in "$@"; do
127 require_relative_inside "$shared_file"
128 done
130 local padded_index
131 padded_index=$(printf "%03d" "$exp_index")
132 local worktree_name="optimize-${spec_name}-exp-${padded_index}"
133 local branch_name
134 branch_name=$(experiment_branch_name "$spec_name" "$padded_index")
135 local worktree_path="$WORKTREE_DIR/$worktree_name"
137 # Check if worktree already exists
138 if [[ -d "$worktree_path" ]]; then
139 if ! git -C "$worktree_path" rev-parse --is-inside-work-tree >/dev/null 2>&1 || \
140 ! is_registered_worktree "$worktree_path"; then
141 echo -e "${RED}Error: Existing path is not a valid registered git worktree: $worktree_path${NC}" >&2
142 echo -e "${RED}Remove or repair that directory before rerunning the experiment.${NC}" >&2
143 return 1
144 fi
146 echo -e "${YELLOW}Worktree already exists: $worktree_path${NC}" >&2
147 reset_worktree_to_base "$worktree_path" "$branch_name" "$base_branch"
148 else
149 mkdir -p "$WORKTREE_DIR"
150 ensure_worktree_exclude
152 # Create worktree from the base branch
153 if ! git worktree add -b "$branch_name" "$worktree_path" "$base_branch" --quiet 2>/dev/null; then
154 if git show-ref --verify --quiet "refs/heads/$branch_name"; then
155 if is_branch_checked_out "$branch_name"; then
156 echo -e "${RED}Error: Existing experiment branch is already checked out: $branch_name${NC}" >&2
157 echo -e "${RED}Clean up the stale worktree before rerunning this experiment.${NC}" >&2
158 return 1
159 fi
161 echo -e "${YELLOW}Resetting existing experiment branch to base: $branch_name -> $base_branch${NC}" >&2
162 git branch -f "$branch_name" "$base_branch" >/dev/null
163 git worktree add "$worktree_path" "$branch_name" --quiet
164 else
165 echo -e "${RED}Error: Failed to create worktree for $branch_name from $base_branch${NC}" >&2
166 return 1
167 fi
168 fi
169 fi
171 # Copy .env files from main repo
172 for f in "$GIT_ROOT"/.env*; do
173 if [[ -f "$f" ]]; then
174 local basename
175 basename=$(basename "$f")
176 if [[ "$basename" != ".env.example" ]]; then
177 cp "$f" "$worktree_path/$basename"
178 fi
179 fi
180 done
182 # Copy shared files
183 for shared_file in "$@"; do
184 if [[ -f "$GIT_ROOT/$shared_file" ]]; then
185 local dir
186 dir=$(dirname "$worktree_path/$shared_file")
187 mkdir -p "$dir"
188 cp "$GIT_ROOT/$shared_file" "$worktree_path/$shared_file"
189 elif [[ -d "$GIT_ROOT/$shared_file" ]]; then
190 local dir
191 dir=$(dirname "$worktree_path/$shared_file")
192 mkdir -p "$dir"
193 rm -rf "$worktree_path/$shared_file"
194 cp -R "$GIT_ROOT/$shared_file" "$worktree_path/$shared_file"
195 fi
196 done
198 echo "$worktree_path"
199}
201# Clean up a single experiment worktree
202cleanup_worktree() {
203 local spec_name="${1:?Error: spec_name required}"
204 local exp_index="${2:?Error: exp_index required}"
205 require_spec_name "$spec_name"
207 local padded_index
208 padded_index=$(printf "%03d" "$exp_index")
209 local worktree_name="optimize-${spec_name}-exp-${padded_index}"
210 local branch_name
211 branch_name=$(experiment_branch_name "$spec_name" "$padded_index")
212 local worktree_path="$WORKTREE_DIR/$worktree_name"
214 if [[ -d "$worktree_path" ]]; then
215 git worktree remove "$worktree_path" --force 2>/dev/null || {
216 # If worktree remove fails, try manual cleanup
217 remove_experiment_dir "$worktree_path" 2>/dev/null || true
218 git worktree prune 2>/dev/null || true
219 }
220 fi
222 # Delete the experiment branch
223 git branch -D "$branch_name" 2>/dev/null || true
225 echo -e "${GREEN}Cleaned up: $worktree_name${NC}" >&2
226}
228# Clean up all experiment worktrees for a spec
229cleanup_all() {
230 local spec_name="${1:?Error: spec_name required}"
231 require_spec_name "$spec_name"
232 local prefix="optimize-${spec_name}-exp-"
233 local count=0
235 if [[ ! -d "$WORKTREE_DIR" ]]; then
236 echo -e "${YELLOW}No worktrees directory found${NC}" >&2
237 return 0
238 fi
240 for worktree_path in "$WORKTREE_DIR"/${prefix}*; do
241 if [[ -d "$worktree_path" ]]; then
242 local worktree_name
243 worktree_name=$(basename "$worktree_path")
244 # Extract index from name
245 local index_str="${worktree_name#$prefix}"
247 git worktree remove "$worktree_path" --force 2>/dev/null || {
248 remove_experiment_dir "$worktree_path" 2>/dev/null || true
249 }
251 # Delete the branch
252 local branch_name
253 branch_name=$(experiment_branch_name "$spec_name" "$index_str")
254 git branch -D "$branch_name" 2>/dev/null || true
256 count=$((count + 1))
257 fi
258 done
260 git worktree prune 2>/dev/null || true
262 # Clean up empty worktree directory
263 if [[ -d "$WORKTREE_DIR" ]] && [[ -z "$(ls -A "$WORKTREE_DIR" 2>/dev/null)" ]]; then
264 rmdir "$WORKTREE_DIR" 2>/dev/null || true
265 fi
267 echo -e "${GREEN}Cleaned up $count experiment worktree(s) for $spec_name${NC}" >&2
268}
270# Count total worktrees (for budget check)
271count_worktrees() {
272 local count=0
273 if [[ -d "$WORKTREE_DIR" ]]; then
274 for worktree_path in "$WORKTREE_DIR"/*; do
275 if [[ -d "$worktree_path" ]] && [[ -e "$worktree_path/.git" ]]; then
276 count=$((count + 1))
277 fi
278 done
279 fi
280 echo "$count"
281}
283# Main
284main() {
285 local command="${1:-help}"
287 case "$command" in
288 create)
289 shift
290 create_worktree "$@"
291 ;;
292 cleanup)
293 shift
294 cleanup_worktree "$@"
295 ;;
296 cleanup-all)
297 shift
298 cleanup_all "$@"
299 ;;
300 count)
301 count_worktrees
302 ;;
303 help)
304 cat << 'EOF'
305Experiment Worktree Manager
307Usage:
308 experiment-worktree.sh create <spec_name> <exp_index> <base_branch> [shared_file ...]
309 experiment-worktree.sh cleanup <spec_name> <exp_index>
310 experiment-worktree.sh cleanup-all <spec_name>
311 experiment-worktree.sh count
313Commands:
314 create Create an experiment worktree with copied shared files
315 cleanup Remove a single experiment worktree and its branch
316 cleanup-all Remove all experiment worktrees for a spec
317 count Count total active worktrees (for budget checking)
319Worktrees: .worktrees/optimize-<spec>-exp-<NNN>/
320Branches: optimize-exp/<spec>/exp-<NNN>
321EOF
322 ;;
323 *)
324 echo -e "${RED}Unknown command: $command${NC}" >&2
325 exit 1
326 ;;
327 esac
328}
330main "$@"