[FEAT][LSP] Worktrees

codeflash/api/aiservice.py

@@ -10,8 +10,9 @@
 from pydantic.json import pydantic_encoder
 
 from codeflash.cli_cmds.console import console, logger
-from codeflash.code_utils.env_utils import get_codeflash_api_key, is_LSP_enabled
+from codeflash.code_utils.env_utils import get_codeflash_api_key
 from codeflash.code_utils.git_utils import get_last_commit_author_if_pr_exists, get_repo_owner_and_name
+from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
 from codeflash.models.models import AIServiceRefinerRequest, CodeStringsMarkdown, OptimizedCandidate
 from codeflash.telemetry.posthog_cf import ph

codeflash/api/cfapi.py

@@ -16,6 +16,7 @@
 from codeflash.code_utils.env_utils import ensure_codeflash_api_key, get_codeflash_api_key, get_pr_number
 from codeflash.code_utils.git_utils import get_current_branch, get_repo_owner_and_name, git_root_dir
 from codeflash.github.PrComment import FileDiffContent, PrComment
+from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.version import __version__
 
 if TYPE_CHECKING:
@@ -101,7 +102,7 @@ def get_user_id() -> Optional[str]:
             if min_version and version.parse(min_version) > version.parse(__version__):
                 msg = "Your Codeflash CLI version is outdated. Please update to the latest version using `pip install --upgrade codeflash`."
                 console.print(f"[bold red]{msg}[/bold red]")
-                if console.quiet:  # lsp
+                if is_LSP_enabled():
                     logger.debug(msg)
                     return f"Error: {msg}"
                 sys.exit(1)

codeflash/cli_cmds/cli.py

@@ -94,6 +94,7 @@ def parse_args() -> Namespace:
         help="Path to the directory of the project, where all the pytest-benchmark tests are located.",
     )
     parser.add_argument("--no-draft", default=False, action="store_true", help="Skip optimization for draft PRs")
+    parser.add_argument("--worktree", default=False, action="store_true", help="Use worktree for optimization")
 
     args, unknown_args = parser.parse_known_args()
     sys.argv[:] = [sys.argv[0], *unknown_args]

codeflash/cli_cmds/console.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import os
 from contextlib import contextmanager
 from itertools import cycle
 from typing import TYPE_CHECKING
@@ -28,6 +29,10 @@
 DEBUG_MODE = logging.getLogger().getEffectiveLevel() == logging.DEBUG
 
 console = Console()
+
+if os.getenv("CODEFLASH_LSP"):
+    console.quiet = True
+
 logging.basicConfig(
     level=logging.INFO,
     handlers=[RichHandler(rich_tracebacks=True, markup=False, console=console, show_path=False, show_time=False)],

codeflash/code_utils/coverage_utils.py

@@ -39,16 +39,20 @@ def build_fully_qualified_name(function_name: str, code_context: CodeOptimizatio
     return full_name
 
 
-def generate_candidates(source_code_path: Path) -> list[str]:
+def generate_candidates(source_code_path: Path) -> set[str]:
     """Generate all the possible candidates for coverage data based on the source code path."""
-    candidates = [source_code_path.name]
+    candidates = set()
+    candidates.add(source_code_path.name)
     current_path = source_code_path.parent
 
+    last_added = source_code_path.name
     while current_path != current_path.parent:
-        candidate_path = str(Path(current_path.name) / candidates[-1])
-        candidates.append(candidate_path)
+        candidate_path = str(Path(current_path.name) / last_added)
+        candidates.add(candidate_path)
+        last_added = candidate_path
         current_path = current_path.parent
 
+    candidates.add(str(source_code_path))
     return candidates
 
 

codeflash/code_utils/env_utils.py

@@ -7,10 +7,11 @@
 from pathlib import Path
 from typing import Any, Optional
 
-from codeflash.cli_cmds.console import console, logger
+from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import exit_with_message
 from codeflash.code_utils.formatter import format_code
 from codeflash.code_utils.shell_utils import read_api_key_from_shell_config
+from codeflash.lsp.helpers import is_LSP_enabled
 
 
 def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool = True) -> bool:  # noqa
@@ -34,11 +35,12 @@ def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool =
 
 @lru_cache(maxsize=1)
 def get_codeflash_api_key() -> str:
-    if console.quiet:  # lsp
-        # prefer shell config over env var in lsp mode
-        api_key = read_api_key_from_shell_config()
-    else:
-        api_key = os.environ.get("CODEFLASH_API_KEY") or read_api_key_from_shell_config()
+    # prefer shell config over env var in lsp mode
+    api_key = (
+        read_api_key_from_shell_config()
+        if is_LSP_enabled()
+        else os.environ.get("CODEFLASH_API_KEY") or read_api_key_from_shell_config()
+    )
 
     api_secret_docs_message = "For more information, refer to the documentation at [https://docs.codeflash.ai/getting-started/codeflash-github-actions#add-your-api-key-to-your-repository-secrets]."  # noqa
     if not api_key:
@@ -125,11 +127,6 @@ def is_ci() -> bool:
     return bool(os.environ.get("CI") or os.environ.get("GITHUB_ACTIONS"))
 
 
-@lru_cache(maxsize=1)
-def is_LSP_enabled() -> bool:
-    return console.quiet
-
-
 def is_pr_draft() -> bool:
     """Check if the PR is draft. in the github action context."""
     event = get_cached_gh_event_data()

codeflash/code_utils/formatter.py

@@ -13,6 +13,7 @@
 import isort
 
 from codeflash.cli_cmds.console import console, logger
+from codeflash.lsp.helpers import is_LSP_enabled
 
 
 def generate_unified_diff(original: str, modified: str, from_file: str, to_file: str) -> str:
@@ -109,8 +110,7 @@ def format_code(
     print_status: bool = True,  # noqa
     exit_on_failure: bool = True,  # noqa
 ) -> str:
-    if console.quiet:
-        # lsp mode
+    if is_LSP_enabled():
         exit_on_failure = False
     with tempfile.TemporaryDirectory() as test_dir_str:
         if isinstance(path, str):

codeflash/code_utils/git_utils.py

@@ -9,14 +9,16 @@
 from functools import cache
 from io import StringIO
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 import git
 from rich.prompt import Confirm
 from unidiff import PatchSet
 
 from codeflash.cli_cmds.console import logger
+from codeflash.code_utils.compat import codeflash_cache_dir
 from codeflash.code_utils.config_consts import N_CANDIDATES
+from codeflash.lsp.helpers import is_LSP_enabled
 
 if TYPE_CHECKING:
     from git import Repo
@@ -192,3 +194,84 @@ def get_last_commit_author_if_pr_exists(repo: Repo | None = None) -> str | None:
         return None
     else:
         return last_commit.author.name
+
+
+worktree_dirs = codeflash_cache_dir / "worktrees"
+patches_dir = codeflash_cache_dir / "patches"
+
+
+def create_worktree_snapshot_commit(worktree_dir: Path, commit_message: str) -> None:
+    repository = git.Repo(worktree_dir, search_parent_directories=True)
+    repository.git.commit("-am", commit_message, "--no-verify")
+
+
+def create_detached_worktree(module_root: Path) -> Optional[Path]:
+    if not check_running_in_git_repo(module_root):
+        logger.warning("Module is not in a git repository. Skipping worktree creation.")
+        return None
+    git_root = git_root_dir()
+    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,
+        stdout=subprocess.DEVNULL if is_LSP_enabled() else None,
+        stderr=subprocess.DEVNULL if is_LSP_enabled() else None,
+    )
+    if result.returncode != 0:
+        logger.error(f"Failed to create worktree: {result.stderr}")
+        return None
+
+    # Get uncommitted diff from the original repo
+    repository = git.Repo(module_root, search_parent_directories=True)
+    repository.git.add("-N", ".")  # add the index for untracked files to be included in the diff
+    uni_diff_text = repository.git.diff(None, "HEAD", ignore_blank_lines=True, ignore_space_at_eol=True)
+
+    if not uni_diff_text.strip():
+        logger.info("No uncommitted changes to copy to worktree.")
+        return worktree_dir
+
+    # Write the diff to a temporary file
+    with tempfile.NamedTemporaryFile(mode="w+", suffix=".codeflash.patch", delete=False) as tmp_patch_file:
+        tmp_patch_file.write(uni_diff_text + "\n")  # the new line here is a must otherwise the last hunk won't be valid
+        tmp_patch_file.flush()
+
+        patch_path = Path(tmp_patch_file.name).resolve()
+
+        # Apply the patch inside the worktree
+        try:
+            subprocess.run(["git", "apply", patch_path], cwd=worktree_dir, check=True)
+            create_worktree_snapshot_commit(worktree_dir, "Initial Snapshot")
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Failed to apply patch to worktree: {e}")
+
+        return worktree_dir
+
+
+def remove_worktree(worktree_dir: Path) -> None:
+    try:
+        repository = git.Repo(worktree_dir, search_parent_directories=True)
+        repository.git.worktree("remove", "--force", worktree_dir)
+    except Exception:
+        logger.exception(f"Failed to remove worktree: {worktree_dir}")
+
+
+def create_diff_patch_from_worktree(worktree_dir: Path, files: list[str], fto_name: str) -> Path:
+    repository = git.Repo(worktree_dir, search_parent_directories=True)
+    uni_diff_text = repository.git.diff(None, "HEAD", *files, ignore_blank_lines=True, ignore_space_at_eol=True)
+
+    if not uni_diff_text:
+        logger.warning("No changes found in worktree.")
+        return None
+
+    if not uni_diff_text.endswith("\n"):
+        uni_diff_text += "\n"
+
+    # write to patches_dir
+    patches_dir.mkdir(parents=True, exist_ok=True)
+    patch_path = patches_dir / f"{worktree_dir.name}.{fto_name}.patch"
+    with patch_path.open("w", encoding="utf8") as f:
+        f.write(uni_diff_text)
+    return patch_path

codeflash/lsp/__init__.py

@@ -1,4 +0,0 @@
-# Silence the console module to prevent stdout pollution
-from codeflash.cli_cmds.console import console
-
-console.quiet = True