worktrees changes

Author: mohammedahmed18

codeflash/code_utils/git_utils.py

@@ -15,7 +15,7 @@
 from rich.prompt import Confirm
 from unidiff import PatchSet
 
-from codeflash.cli_cmds.console import logger
+from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.compat import codeflash_cache_dir
 from codeflash.code_utils.config_consts import N_CANDIDATES
 
@@ -212,13 +212,17 @@ def create_detached_worktree(module_root: Path) -> Optional[Path]:
     current_time_str = time.strftime("%Y%m%d-%H%M%S")
     worktree_dir = worktree_dirs / f"{git_root.name}-{current_time_str}"
 
-    result = subprocess.run(["git", "worktree", "add", "-d", str(worktree_dir)], cwd=git_root, check=True)
+    result = subprocess.run(
+        ["git", "worktree", "add", "-d", str(worktree_dir)],
+        cwd=git_root,
+        check=True,
+        stdout=subprocess.DEVNULL if console.quiet else None,
+        stderr=subprocess.DEVNULL if console.quiet else None,
+    )
     if result.returncode != 0:
         logger.error(f"Failed to create worktree: {result.stderr}")
         return None
 
-    print(result.stdout)
-
     # Get uncommitted diff from the original repo
     repository = git.Repo(module_root, search_parent_directories=True)
     uni_diff_text = repository.git.diff(None, "HEAD", ignore_blank_lines=True, ignore_space_at_eol=True)

codeflash/lsp/beta.py

@@ -9,6 +9,7 @@
 from pygls import uris
 
 from codeflash.api.cfapi import get_codeflash_api_key, get_user_id
+from codeflash.code_utils.git_utils import create_diff_from_worktree
 from codeflash.code_utils.shell_utils import save_api_key_to_rc
 from codeflash.either import is_successful
 from codeflash.lsp.server import CodeflashLanguageServer, CodeflashLanguageServerProtocol
@@ -331,17 +332,23 @@ def perform_function_optimization(  # noqa: PLR0911
             "message": f"No best optimizations found for function {function_to_optimize_qualified_name}",
         }
 
+    # generate a patch for the optimization
+    relative_file_paths = [code_string.file_path for code_string in code_context.read_writable_code.code_strings]
+    patch_path = create_diff_from_worktree(
+        server.optimizer.current_worktree,
+        relative_file_paths,
+        server.optimizer.current_function_optimizer.function_to_optimize.qualified_name,
+    )
+
     optimized_source = best_optimization.candidate.source_code.markdown
     speedup = original_code_baseline.runtime / best_optimization.runtime
 
     server.show_message_log(f"Optimization completed for {params.functionName} with {speedup:.2f}x speedup", "Info")
-    diff_patch_files = server.optimizer.patch_files
 
     # CRITICAL: Clear the function filter after optimization to prevent state corruption
+    server.optimizer.cleanup_temporary_paths()
     server.optimizer.args.function = None
-    server.optimizer.patch_files = []
     server.optimizer.current_worktree = None
-    server.optimizer.cleanup_temporary_paths()
     server.show_message_log("Cleared function filter to prevent state corruption", "Info")
 
     return {
@@ -350,5 +357,5 @@ def perform_function_optimization(  # noqa: PLR0911
         "message": "Optimization completed successfully",
         "extra": f"Speedup: {speedup:.2f}x faster",
         "optimization": optimized_source,
-        "diff_patch_files": diff_patch_files,
+        "patch_path": patch_path,
     }

codeflash/lsp/server.py

@@ -55,6 +55,7 @@ def prepare_optimizer_arguments(self, config_file: Path) -> None:
         args = parse_args()
         args.config_file = config_file
         args.no_pr = True  # LSP server should not create PRs
+        args.worktree = True
         args = process_pyproject_config(args)
         self.args = args
         # avoid initializing the optimizer during initialization, because it can cause an error if the api key is invalid

codeflash/optimization/function_optimizer.py

@@ -1252,6 +1252,10 @@ def process_review(
                 trace_id=self.function_trace_id, is_optimization_found=best_optimization is not None
             )
 
+        # If worktree mode, do not revert code and helpers, we create snapshot commit after each function optimization
+        if self.args.worktree:
+            return
+
         if raise_pr and (
             self.args.all
             or env_utils.get_pr_number()

codeflash/optimization/optimizer.py

@@ -414,16 +414,17 @@ def cleanup_replay_tests(self) -> None:
             cleanup_paths([self.replay_tests_dir])
 
     def cleanup_temporary_paths(self) -> None:
-        if self.current_function_optimizer:
-            self.current_function_optimizer.cleanup_generated_files()
-
         if hasattr(get_run_tmp_file, "tmpdir"):
             get_run_tmp_file.tmpdir.cleanup()
             del get_run_tmp_file.tmpdir
-        paths_to_clean = [self.test_cfg.concolic_test_root_dir, self.replay_tests_dir]
+
         if self.current_worktree:
             remove_worktree(self.current_worktree)
-        cleanup_paths(paths_to_clean)
+            return
+
+        if self.current_function_optimizer:
+            self.current_function_optimizer.cleanup_generated_files()
+        cleanup_paths([self.test_cfg.concolic_test_root_dir, self.replay_tests_dir])
 
     def worktree_mode(self) -> None:
         if self.current_worktree: