fixes

Author: mohammedahmed18

codeflash/api/cfapi.py

@@ -13,7 +13,12 @@
 from pydantic.json import pydantic_encoder
 
 from codeflash.cli_cmds.console import console, logger
-from codeflash.code_utils.env_utils import ensure_codeflash_api_key, get_codeflash_api_key, get_pr_number
+from codeflash.code_utils.env_utils import (
+    ensure_codeflash_api_key,
+    get_codeflash_api_key,
+    get_pr_number,
+    is_LSP_enabled,
+)
 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.version import __version__
@@ -101,7 +106,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/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

@@ -15,9 +15,10 @@
 from rich.prompt import Confirm
 from unidiff import PatchSet
 
-from codeflash.cli_cmds.console import console, logger
+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
@@ -216,15 +217,16 @@ def create_detached_worktree(module_root: Path) -> Optional[Path]:
         ["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,
+        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():

codeflash/lsp/beta.py

@@ -46,13 +46,7 @@ def get_optimizable_functions(
     file_path = Path(uris.to_fs_path(params.textDocument.uri))
     server.show_message_log(f"Getting optimizable functions for: {file_path}", "Info")
 
-    server.optimizer.worktree_mode()
-
-    args = server.optimizer.args
-
-    original_relative_file_path = file_path.relative_to(args.base_project_root)
-
-    server.optimizer.args.file = server.optimizer.current_worktree / original_relative_file_path
+    server.optimizer.args.file = file_path
     server.optimizer.args.function = None  # Always get ALL functions, not just one
     server.optimizer.args.previous_checkpoint_functions = False
 
@@ -75,11 +69,15 @@ def initialize_function_optimization(
 ) -> dict[str, str]:
     file_path = Path(uris.to_fs_path(params.textDocument.uri))
     server.show_message_log(f"Initializing optimization for function: {params.functionName} in {file_path}", "Info")
+    if server.optimizer is None:
+        _initialize_optimizer_if_valid(server)
     server.optimizer.worktree_mode()
+    original_args, _ = server.optimizer.original_args_and_test_cfg
 
     server.optimizer.args.function = params.functionName
-    original_relative_file_path = file_path.relative_to(server.optimizer.args.base_project_root)
+    original_relative_file_path = file_path.relative_to(original_args.project_root)
     server.optimizer.args.file = server.optimizer.current_worktree / original_relative_file_path
+    server.optimizer.args.previous_checkpoint_functions = False
 
     server.show_message_log(
         f"Args set - function: {server.optimizer.args.function}, file: {server.optimizer.args.file}", "Info"
@@ -338,10 +336,14 @@ def perform_function_optimization(  # noqa: PLR0911
             "patch_path": patch_path,
         }
     finally:
-        server.optimizer.cleanup_temporary_paths()
-        # restore args and test cfg
-        if server.optimizer.original_args_and_test_cfg:
-            server.optimizer.args, server.optimizer.test_cfg = server.optimizer.original_args_and_test_cfg
-        server.optimizer.args.function = None
-        server.optimizer.current_worktree = None
-        server.optimizer.current_function_optimizer = None
+        cleanup_the_optimizer(server)
+
+
+def cleanup_the_optimizer(server: CodeflashLanguageServer) -> None:
+    server.optimizer.cleanup_temporary_paths()
+    # restore args and test cfg
+    if server.optimizer.original_args_and_test_cfg:
+        server.optimizer.args, server.optimizer.test_cfg = server.optimizer.original_args_and_test_cfg
+    server.optimizer.args.function = None
+    server.optimizer.current_worktree = None
+    server.optimizer.current_function_optimizer = None

codeflash/lsp/helpers.py

@@ -0,0 +1,8 @@
+from functools import lru_cache
+
+from codeflash.cli_cmds.console import console
+
+
+@lru_cache(maxsize=1)
+def is_LSP_enabled() -> bool:
+    return console.quiet

codeflash/optimization/optimizer.py

@@ -446,7 +446,6 @@ def mutate_args_for_worktree_mode(self, worktree_dir: Path) -> None:
         self.original_args_and_test_cfg = (saved_args, saved_test_cfg)
 
         project_root = self.args.project_root
-        self.args.base_project_root = project_root  # keep a reference to the original project root
         module_root = self.args.module_root
         relative_module_root = module_root.relative_to(project_root)
         relative_optimized_file = self.args.file.relative_to(project_root) if self.args.file else None

codeflash/verification/concolic_testing.py

@@ -10,9 +10,9 @@
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
 from codeflash.code_utils.concolic_utils import clean_concolic_tests
-from codeflash.code_utils.env_utils import is_LSP_enabled
 from codeflash.code_utils.static_analysis import has_typed_parameters
 from codeflash.discovery.discover_unit_tests import discover_unit_tests
+from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.telemetry.posthog_cf import ph
 from codeflash.verification.verification_utils import TestConfig