Merge pull request #1522 from codeflash-ai/cf-fix-worktree-module-root-path

Fix worktree mode filtering out all diff-discovered functions

Author: KRRT7

codeflash/optimization/function_optimizer.py

@@ -13,6 +13,7 @@
 from typing import TYPE_CHECKING, Callable
 
 import libcst as cst
+from git import Repo as GitRepo
 from rich.console import Group
 from rich.panel import Panel
 from rich.syntax import Syntax
@@ -2222,11 +2223,11 @@ def process_review(
                 console.print(Panel(panel_content, title="Optimization Review", border_style=display_info[1]))
 
         if raise_pr or staging_review:
-            data["root_dir"] = git_root_dir()
+            data["root_dir"] = git_root_dir(GitRepo(str(self.args.module_root), search_parent_directories=True))
         if raise_pr and not staging_review and opt_review_result.review != "low":
             # Ensure root_dir is set for PR creation (needed for async functions that skip opt_review)
             if "root_dir" not in data:
-                data["root_dir"] = git_root_dir()
+                data["root_dir"] = git_root_dir(GitRepo(str(self.args.module_root), search_parent_directories=True))
             data["git_remote"] = self.args.git_remote
             # Remove language from data dict as check_create_pr doesn't accept it
             pr_data = {k: v for k, v in data.items() if k != "language"}

codeflash/optimization/optimizer.py

@@ -183,18 +183,54 @@ def get_optimizable_functions(self) -> tuple[dict[Path, list[FunctionToOptimize]
         """Discover functions to optimize."""
         from codeflash.discovery.functions_to_optimize import get_functions_to_optimize
 
-        return get_functions_to_optimize(
+        # In worktree mode for git-diff discovery, file paths come from the original repo
+        # (via get_git_diff using cwd), but module_root/project_root have been mirrored to
+        # the worktree. Use the original roots for filtering so path comparisons match,
+        # then remap the discovered file paths to the worktree.
+        project_root = self.args.project_root
+        module_root = self.args.module_root
+        use_original_roots = (
+            self.current_worktree and self.original_args_and_test_cfg and not self.args.all and not self.args.file
+        )
+        if use_original_roots:
+            assert self.original_args_and_test_cfg is not None
+            original_args, _ = self.original_args_and_test_cfg
+            project_root = original_args.project_root
+            module_root = original_args.module_root
+
+        result = get_functions_to_optimize(
             optimize_all=self.args.all,
             replay_test=self.args.replay_test,
             file=self.args.file,
             only_get_this_function=self.args.function,
             test_cfg=self.test_cfg,
             ignore_paths=self.args.ignore_paths,
-            project_root=self.args.project_root,
-            module_root=self.args.module_root,
+            project_root=project_root,
+            module_root=module_root,
             previous_checkpoint_functions=self.args.previous_checkpoint_functions,
         )
 
+        # Remap discovered file paths from the original repo to the worktree so
+        # downstream optimization reads/writes happen in the worktree.
+        if use_original_roots:
+            import dataclasses
+
+            assert self.current_worktree is not None
+            original_git_root = git_root_dir()
+            file_to_funcs, count, trace = result
+            remapped: dict[Path, list[FunctionToOptimize]] = {}
+            for file_path, funcs in file_to_funcs.items():
+                new_path = mirror_path(Path(file_path), original_git_root, self.current_worktree)
+                remapped[new_path] = [
+                    dataclasses.replace(
+                        func, file_path=mirror_path(func.file_path, original_git_root, self.current_worktree)
+                    )
+                    for func in funcs
+                ]
+            return remapped, count, trace
+
+        return result
+
     def create_function_optimizer(
         self,
         function_to_optimize: FunctionToOptimize,