From 3f74dd177a35e4989c2ce6613b927cec7ef26d8b Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Wed, 11 Mar 2026 12:58:59 -0700 Subject: [PATCH 1/4] fix: ensure worktree cleanup on SIGTERM, SIGHUP, and atexit Register signal handlers for SIGTERM and SIGHUP in run_with_args so that the git worktree is removed when the process is killed or the terminal is closed. An atexit handler provides a last-resort safety net for other unexpected exit paths. Co-Authored-By: Claude Opus 4.6 --- codeflash/optimization/optimizer.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 21fe83ff2..954c424ba 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -735,7 +735,27 @@ def mirror_path(path: Path, src_root: Path, dest_root: Path) -> Path: def run_with_args(args: Namespace) -> None: + import atexit + import signal + optimizer = None + original_sigterm = signal.getsignal(signal.SIGTERM) + original_sighup = signal.getsignal(signal.SIGHUP) + + def cleanup_worktree_on_exit() -> None: + if optimizer and optimizer.current_worktree: + remove_worktree(optimizer.current_worktree) + + def signal_handler(signum: int, frame: object) -> None: + logger.warning(f"Signal {signum} received. Cleaning up worktree and exiting…") + if optimizer: + optimizer.cleanup_temporary_paths() + raise SystemExit(128 + signum) + + atexit.register(cleanup_worktree_on_exit) + signal.signal(signal.SIGTERM, signal_handler) + signal.signal(signal.SIGHUP, signal_handler) + try: optimizer = Optimizer(args) optimizer.run() @@ -745,3 +765,7 @@ def run_with_args(args: Namespace) -> None: optimizer.cleanup_temporary_paths() raise SystemExit from None + finally: + atexit.unregister(cleanup_worktree_on_exit) + signal.signal(signal.SIGTERM, original_sigterm) + signal.signal(signal.SIGHUP, original_sighup) From b9cbec185ced1e397dedd8d8491b425027aa9f30 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Wed, 11 Mar 2026 13:06:25 -0700 Subject: [PATCH 2/4] fix: add SIGQUIT handler for worktree cleanup Co-Authored-By: Claude Opus 4.6 --- codeflash/optimization/optimizer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 954c424ba..18523f4f0 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -741,6 +741,7 @@ def run_with_args(args: Namespace) -> None: optimizer = None original_sigterm = signal.getsignal(signal.SIGTERM) original_sighup = signal.getsignal(signal.SIGHUP) + original_sigquit = signal.getsignal(signal.SIGQUIT) def cleanup_worktree_on_exit() -> None: if optimizer and optimizer.current_worktree: @@ -755,6 +756,7 @@ def signal_handler(signum: int, frame: object) -> None: atexit.register(cleanup_worktree_on_exit) signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGHUP, signal_handler) + signal.signal(signal.SIGQUIT, signal_handler) try: optimizer = Optimizer(args) @@ -769,3 +771,4 @@ def signal_handler(signum: int, frame: object) -> None: atexit.unregister(cleanup_worktree_on_exit) signal.signal(signal.SIGTERM, original_sigterm) signal.signal(signal.SIGHUP, original_sighup) + signal.signal(signal.SIGQUIT, original_sigquit) From 0e8befa602aa50b855c28734bc420dce136c7427 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Wed, 11 Mar 2026 13:07:04 -0700 Subject: [PATCH 3/4] fix: add SIGPIPE handler for worktree cleanup Co-Authored-By: Claude Opus 4.6 --- codeflash/optimization/optimizer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 18523f4f0..277dcb5c6 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -742,6 +742,7 @@ def run_with_args(args: Namespace) -> None: original_sigterm = signal.getsignal(signal.SIGTERM) original_sighup = signal.getsignal(signal.SIGHUP) original_sigquit = signal.getsignal(signal.SIGQUIT) + original_sigpipe = signal.getsignal(signal.SIGPIPE) def cleanup_worktree_on_exit() -> None: if optimizer and optimizer.current_worktree: @@ -757,6 +758,7 @@ def signal_handler(signum: int, frame: object) -> None: signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGHUP, signal_handler) signal.signal(signal.SIGQUIT, signal_handler) + signal.signal(signal.SIGPIPE, signal_handler) try: optimizer = Optimizer(args) @@ -772,3 +774,4 @@ def signal_handler(signum: int, frame: object) -> None: signal.signal(signal.SIGTERM, original_sigterm) signal.signal(signal.SIGHUP, original_sighup) signal.signal(signal.SIGQUIT, original_sigquit) + signal.signal(signal.SIGPIPE, original_sigpipe) From 9edb21deecd649de037922102dbeb6bca10068ac Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Wed, 11 Mar 2026 13:10:51 -0700 Subject: [PATCH 4/4] fix: clean up stale worktrees at startup after SIGKILL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Write a .codeflash.pid file when creating a worktree. On startup, scan the worktrees directory and remove any whose owning process is no longer alive — handles the case where a previous run was killed with SIGKILL. Co-Authored-By: Claude Opus 4.6 --- codeflash/code_utils/git_worktree_utils.py | 35 ++++++++++++++++++++++ codeflash/optimization/optimizer.py | 3 ++ 2 files changed, 38 insertions(+) diff --git a/codeflash/code_utils/git_worktree_utils.py b/codeflash/code_utils/git_worktree_utils.py index 3dcba708e..1e7738763 100644 --- a/codeflash/code_utils/git_worktree_utils.py +++ b/codeflash/code_utils/git_worktree_utils.py @@ -1,6 +1,7 @@ from __future__ import annotations import configparser +import os import shutil import stat import subprocess @@ -65,6 +66,10 @@ def create_detached_worktree(module_root: Path) -> Optional[Path]: repository.git.worktree("add", "-d", str(worktree_dir)) + # Write PID file so stale worktrees can be detected after SIGKILL + pid_file = worktree_dir / ".codeflash.pid" + pid_file.write_text(str(os.getpid()), encoding="utf-8") + # Get uncommitted diff from the original repo repository.git.add("-N", ".") # add the index for untracked files to be included in the diff exclude_binary_files = [":!*.pyc", ":!*.pyo", ":!*.pyd", ":!*.so", ":!*.dll", ":!*.whl", ":!*.egg", ":!*.egg-info", ":!*.pyz", ":!*.pkl", ":!*.pickle", ":!*.joblib", ":!*.npy", ":!*.npz", ":!*.h5", ":!*.hdf5", ":!*.pth", ":!*.pt", ":!*.pb", ":!*.onnx", ":!*.db", ":!*.sqlite", ":!*.sqlite3", ":!*.feather", ":!*.parquet", ":!*.jpg", ":!*.jpeg", ":!*.png", ":!*.gif", ":!*.bmp", ":!*.tiff", ":!*.webp", ":!*.wav", ":!*.mp3", ":!*.ogg", ":!*.flac", ":!*.mp4", ":!*.avi", ":!*.mov", ":!*.mkv", ":!*.pdf", ":!*.doc", ":!*.docx", ":!*.xls", ":!*.xlsx", ":!*.ppt", ":!*.pptx", ":!*.zip", ":!*.rar", ":!*.tar", ":!*.tar.gz", ":!*.tgz", ":!*.bz2", ":!*.xz"] # fmt: off @@ -119,6 +124,36 @@ def remove_worktree(worktree_dir: Path) -> None: logger.exception(f"Failed to remove worktree: {worktree_dir}") +def is_process_alive(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True # process exists but we can't signal it + return True + + +def cleanup_stale_worktrees() -> None: + """Remove worktrees left behind by killed processes (e.g. SIGKILL).""" + if not worktree_dirs.exists(): + return + for entry in worktree_dirs.iterdir(): + if not entry.is_dir(): + continue + pid_file = entry / ".codeflash.pid" + if pid_file.exists(): + try: + pid = int(pid_file.read_text(encoding="utf-8").strip()) + except (ValueError, OSError): + pid = None + if pid is not None and is_process_alive(pid): + continue # worktree is still in use + # No PID file or owning process is dead — stale worktree + logger.info(f"Removing stale worktree: {entry}") + remove_worktree(entry) + + def create_diff_patch_from_worktree( worktree_dir: Path, files: list[Path], fto_name: Optional[str] = None ) -> Optional[Path]: diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 277dcb5c6..82ae5520a 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -22,6 +22,7 @@ from codeflash.code_utils.env_utils import get_pr_number, is_pr_draft from codeflash.code_utils.git_utils import check_running_in_git_repo, git_root_dir from codeflash.code_utils.git_worktree_utils import ( + cleanup_stale_worktrees, create_detached_worktree, create_diff_patch_from_worktree, create_worktree_snapshot_commit, @@ -738,6 +739,8 @@ def run_with_args(args: Namespace) -> None: import atexit import signal + cleanup_stale_worktrees() + optimizer = None original_sigterm = signal.getsignal(signal.SIGTERM) original_sighup = signal.getsignal(signal.SIGHUP)