Merge branch 'main' into 3.14-in-CI

Author: KRRT7

.github/workflows/e2e-init-optimization.yaml

@@ -19,7 +19,7 @@ jobs:
       COLUMNS: 110
       MAX_RETRIES: 3
       RETRY_DELAY: 5
-      EXPECTED_IMPROVEMENT_PCT: 30
+      EXPECTED_IMPROVEMENT_PCT: 10
       CODEFLASH_END_TO_END: 1
     steps:
       - name: 🛎️ Checkout

.github/workflows/e2e-topological-sort.yaml

@@ -1,4 +1,4 @@
-name: E2E - Topological Sort
+name: E2E - Topological Sort (Worktree)
 
 on:
   pull_request:
@@ -8,7 +8,7 @@ on:
   workflow_dispatch:
 
 jobs:
-  topological-sort-optimization:
+  topological-sort-worktree-optimization:
     # Dynamically determine if environment is needed only when workflow files change and contributor is external
     environment: ${{ (github.event_name == 'workflow_dispatch' || (contains(toJSON(github.event.pull_request.files.*.filename), '.github/workflows/') && github.event.pull_request.user.login != 'misrasaurabh1' && github.event.pull_request.user.login != 'KRRT7')) && 'external-trusted-contributors' || '' }}
     runs-on: ubuntu-latest
@@ -90,4 +90,4 @@ jobs:
       - name: Run Codeflash to optimize code
         id: optimize_code
         run: |
-          uv run python tests/scripts/end_to_end_test_topological_sort.py
+          uv run python tests/scripts/end_to_end_test_topological_sort_worktree.py

codeflash/cli_cmds/cmd_init.py

@@ -155,20 +155,30 @@ def ask_run_end_to_end_test(args: Namespace) -> None:
         run_end_to_end_test(args, bubble_sort_path, bubble_sort_test_path)
 
 
-def is_valid_pyproject_toml(pyproject_toml_path: Path) -> dict[str, Any] | None:
+def is_valid_pyproject_toml(pyproject_toml_path: Path) -> tuple[dict[str, Any] | None, str]:  # noqa: PLR0911
     if not pyproject_toml_path.exists():
-        return None
+        return None, f"Configuration file not found: {pyproject_toml_path}"
+
     try:
         config, _ = parse_config_file(pyproject_toml_path)
-    except Exception:
-        return None
+    except Exception as e:
+        return None, f"Failed to parse configuration: {e}"
+
+    module_root = config.get("module_root")
+    if not module_root:
+        return None, "Missing required field: 'module_root'"
+
+    if not Path(module_root).is_dir():
+        return None, f"Invalid 'module_root': directory does not exist at {module_root}"
+
+    tests_root = config.get("tests_root")
+    if not tests_root:
+        return None, "Missing required field: 'tests_root'"
 
-    if "module_root" not in config or config["module_root"] is None or not Path(config["module_root"]).is_dir():
-        return None
-    if "tests_root" not in config or config["tests_root"] is None or not Path(config["tests_root"]).is_dir():
-        return None
+    if not Path(tests_root).is_dir():
+        return None, f"Invalid 'tests_root': directory does not exist at {tests_root}"
 
-    return config
+    return config, ""
 
 
 def should_modify_pyproject_toml() -> tuple[bool, dict[str, Any] | None]:
@@ -180,7 +190,7 @@ def should_modify_pyproject_toml() -> tuple[bool, dict[str, Any] | None]:
 
     pyproject_toml_path = Path.cwd() / "pyproject.toml"
 
-    config = is_valid_pyproject_toml(pyproject_toml_path)
+    config, _message = is_valid_pyproject_toml(pyproject_toml_path)
     if config is None:
         return True, None
 
@@ -199,9 +209,7 @@ def __init__(self) -> None:
         self.Question.brackets_color = inquirer.themes.term.bright_blue
         self.Question.default_color = inquirer.themes.term.bright_cyan
         self.List.selection_color = inquirer.themes.term.bright_blue
-        self.List.selection_cursor = "⚡"
         self.Checkbox.selection_color = inquirer.themes.term.bright_blue
-        self.Checkbox.selection_cursor = "⚡"
         self.Checkbox.selected_icon = "✅"
         self.Checkbox.unselected_icon = "⬜"
 
@@ -633,7 +641,7 @@ def check_for_toml_or_setup_file() -> str | None:
 
 def install_github_actions(override_formatter_check: bool = False) -> None:  # noqa: FBT001, FBT002
     try:
-        config, config_file_path = parse_config_file(override_formatter_check=override_formatter_check)
+        config, _config_file_path = parse_config_file(override_formatter_check=override_formatter_check)
 
         ph("cli-github-actions-install-started")
         try:

codeflash/code_utils/code_extractor.py

@@ -528,16 +528,29 @@ def add_needed_imports_from_module(
 
     try:
         for mod in gatherer.module_imports:
+            # Skip __future__ imports as they cannot be imported directly
+            # __future__ imports should only be imported with specific objects i.e from __future__ import annotations
+            if mod == "__future__":
+                continue
             if mod not in dotted_import_collector.imports:
                 AddImportsVisitor.add_needed_import(dst_context, mod)
             RemoveImportsVisitor.remove_unused_import(dst_context, mod)
+        aliased_objects = set()
+        for mod, alias_pairs in gatherer.alias_mapping.items():
+            for alias_pair in alias_pairs:
+                if alias_pair[0] and alias_pair[1]:  # Both name and alias exist
+                    aliased_objects.add(f"{mod}.{alias_pair[0]}")
+
         for mod, obj_seq in gatherer.object_mapping.items():
             for obj in obj_seq:
                 if (
                     f"{mod}.{obj}" in helper_functions_fqn or dst_context.full_module_name == mod  # avoid circular deps
                 ):
                     continue  # Skip adding imports for helper functions already in the context
 
+                if f"{mod}.{obj}" in aliased_objects:
+                    continue
+
                 # Handle star imports by resolving them to actual symbol names
                 if obj == "*":
                     resolved_symbols = resolve_star_import(mod, project_root)
@@ -559,6 +572,8 @@ def add_needed_imports_from_module(
         return dst_module_code
 
     for mod, asname in gatherer.module_aliases.items():
+        if not asname:
+            continue
         if f"{mod}.{asname}" not in dotted_import_collector.imports:
             AddImportsVisitor.add_needed_import(dst_context, mod, asname=asname)
         RemoveImportsVisitor.remove_unused_import(dst_context, mod, asname=asname)
@@ -568,12 +583,16 @@ def add_needed_imports_from_module(
             if f"{mod}.{alias_pair[0]}" in helper_functions_fqn:
                 continue
 
+            if not alias_pair[0] or not alias_pair[1]:
+                continue
+
             if f"{mod}.{alias_pair[1]}" not in dotted_import_collector.imports:
                 AddImportsVisitor.add_needed_import(dst_context, mod, alias_pair[0], asname=alias_pair[1])
             RemoveImportsVisitor.remove_unused_import(dst_context, mod, alias_pair[0], asname=alias_pair[1])
 
     try:
-        transformed_module = AddImportsVisitor(dst_context).transform_module(parsed_dst_module)
+        add_imports_visitor = AddImportsVisitor(dst_context)
+        transformed_module = add_imports_visitor.transform_module(parsed_dst_module)
         transformed_module = RemoveImportsVisitor(dst_context).transform_module(transformed_module)
         return transformed_module.code.lstrip("\n")
     except Exception as e:

codeflash/code_utils/git_worktree_utils.py

@@ -34,6 +34,12 @@ def get_git_project_id() -> str:
 
 def create_worktree_snapshot_commit(worktree_dir: Path, commit_message: str) -> None:
     repository = git.Repo(worktree_dir, search_parent_directories=True)
+    with repository.config_writer() as cw:
+        if not cw.has_option("user", "name"):
+            cw.set_value("user", "name", "Codeflash Bot")
+        if not cw.has_option("user", "email"):
+            cw.set_value("user", "email", "bot@codeflash.ai")
+
     repository.git.add(".")
     repository.git.commit("-m", commit_message, "--no-verify")
 

codeflash/lsp/beta.py

@@ -1,26 +1,24 @@
 from __future__ import annotations
 
-import contextlib
+import asyncio
 import os
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
-import git
 from pygls import uris
 
 from codeflash.api.cfapi import get_codeflash_api_key, get_user_id
 from codeflash.cli_cmds.cli import process_pyproject_config
-from codeflash.cli_cmds.console import code_print
 from codeflash.code_utils.git_utils import git_root_dir
-from codeflash.code_utils.git_worktree_utils import create_diff_patch_from_worktree
 from codeflash.code_utils.shell_utils import save_api_key_to_rc
 from codeflash.discovery.functions_to_optimize import (
     filter_functions,
     get_functions_inside_a_commit,
     get_functions_within_git_diff,
 )
 from codeflash.either import is_successful
+from codeflash.lsp.features.perform_optimization import sync_perform_optimization
 from codeflash.lsp.server import CodeflashLanguageServer
 
 if TYPE_CHECKING:
@@ -71,7 +69,6 @@ class OptimizableFunctionsInCommitParams:
     commit_hash: str
 
 
-# server = CodeflashLanguageServer("codeflash-language-server", "v1.0", protocol_cls=CodeflashLanguageServerProtocol)
 server = CodeflashLanguageServer("codeflash-language-server", "v1.0")
 
 
@@ -157,11 +154,13 @@ def _find_pyproject_toml(workspace_path: str) -> tuple[Path | None, bool]:
 
 # should be called the first thing to initialize and validate the project
 @server.feature("initProject")
-def init_project(server: CodeflashLanguageServer, params: ValidateProjectParams) -> dict[str, str]:  # noqa: PLR0911
+def init_project(server: CodeflashLanguageServer, params: ValidateProjectParams) -> dict[str, str]:
     from codeflash.cli_cmds.cmd_init import is_valid_pyproject_toml
 
-    pyproject_toml_path: Path | None = getattr(params, "config_file", None) or getattr(server.args, "config_file", None)
+    # Always process args in the init project, the extension can call
+    server.args_processed_before = False
 
+    pyproject_toml_path: Path | None = getattr(params, "config_file", None) or getattr(server.args, "config_file", None)
     if pyproject_toml_path is not None:
         # if there is a config file provided use it
         server.prepare_optimizer_arguments(pyproject_toml_path)
@@ -192,20 +191,12 @@ def init_project(server: CodeflashLanguageServer, params: ValidateProjectParams)
         }
 
     server.show_message_log("Validating project...", "Info")
-    config = is_valid_pyproject_toml(pyproject_toml_path)
+    config, reason = is_valid_pyproject_toml(pyproject_toml_path)
     if config is None:
         server.show_message_log("pyproject.toml is not valid", "Error")
-        return {"status": "error", "message": "not valid", "pyprojectPath": pyproject_toml_path}
+        return {"status": "error", "message": f"reason: {reason}", "pyprojectPath": pyproject_toml_path}
 
     args = process_args(server)
-    repo = git.Repo(args.module_root, search_parent_directories=True)
-    if repo.bare:
-        return {"status": "error", "message": "Repository is in bare state"}
-
-    try:
-        _ = repo.head.commit
-    except Exception:
-        return {"status": "error", "message": "Repository has no commits (unborn HEAD)"}
 
     return {"status": "success", "moduleRoot": args.module_root, "pyprojectPath": pyproject_toml_path, "root": root}
 
@@ -339,115 +330,15 @@ def initialize_function_optimization(
 
 
 @server.feature("performFunctionOptimization")
-@server.thread()
-def perform_function_optimization(
+async def perform_function_optimization(
     server: CodeflashLanguageServer, params: FunctionOptimizationParams
 ) -> dict[str, str]:
+    loop = asyncio.get_running_loop()
     try:
-        server.show_message_log(f"Starting optimization for function: {params.functionName}", "Info")
-        should_run_experiment, code_context, original_helper_code = server.current_optimization_init_result
-        function_optimizer = server.optimizer.current_function_optimizer
-        current_function = function_optimizer.function_to_optimize
-
-        code_print(
-            code_context.read_writable_code.flat,
-            file_name=current_function.file_path,
-            function_name=current_function.function_name,
-        )
-
-        optimizable_funcs = {current_function.file_path: [current_function]}
-
-        devnull_writer = open(os.devnull, "w")  # noqa
-        with contextlib.redirect_stdout(devnull_writer):
-            function_to_tests, _num_discovered_tests = server.optimizer.discover_tests(optimizable_funcs)
-            function_optimizer.function_to_tests = function_to_tests
-
-        test_setup_result = function_optimizer.generate_and_instrument_tests(
-            code_context, should_run_experiment=should_run_experiment
-        )
-        if not is_successful(test_setup_result):
-            return {"functionName": params.functionName, "status": "error", "message": test_setup_result.failure()}
-        (
-            generated_tests,
-            function_to_concolic_tests,
-            concolic_test_str,
-            optimizations_set,
-            generated_test_paths,
-            generated_perf_test_paths,
-            instrumented_unittests_created_for_function,
-            original_conftest_content,
-        ) = test_setup_result.unwrap()
-
-        baseline_setup_result = function_optimizer.setup_and_establish_baseline(
-            code_context=code_context,
-            original_helper_code=original_helper_code,
-            function_to_concolic_tests=function_to_concolic_tests,
-            generated_test_paths=generated_test_paths,
-            generated_perf_test_paths=generated_perf_test_paths,
-            instrumented_unittests_created_for_function=instrumented_unittests_created_for_function,
-            original_conftest_content=original_conftest_content,
-        )
-
-        if not is_successful(baseline_setup_result):
-            return {"functionName": params.functionName, "status": "error", "message": baseline_setup_result.failure()}
-
-        (
-            function_to_optimize_qualified_name,
-            function_to_all_tests,
-            original_code_baseline,
-            test_functions_to_remove,
-            file_path_to_helper_classes,
-        ) = baseline_setup_result.unwrap()
-
-        best_optimization = function_optimizer.find_and_process_best_optimization(
-            optimizations_set=optimizations_set,
-            code_context=code_context,
-            original_code_baseline=original_code_baseline,
-            original_helper_code=original_helper_code,
-            file_path_to_helper_classes=file_path_to_helper_classes,
-            function_to_optimize_qualified_name=function_to_optimize_qualified_name,
-            function_to_all_tests=function_to_all_tests,
-            generated_tests=generated_tests,
-            test_functions_to_remove=test_functions_to_remove,
-            concolic_test_str=concolic_test_str,
-        )
-
-        if not best_optimization:
-            server.show_message_log(
-                f"No best optimizations found for function {function_to_optimize_qualified_name}", "Warning"
-            )
-            return {
-                "functionName": params.functionName,
-                "status": "error",
-                "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]
-
-        speedup = original_code_baseline.runtime / best_optimization.runtime
-
-        patch_path = create_diff_patch_from_worktree(
-            server.optimizer.current_worktree, relative_file_paths, function_to_optimize_qualified_name
-        )
-
-        if not patch_path:
-            return {
-                "functionName": params.functionName,
-                "status": "error",
-                "message": "Failed to create a patch for optimization",
-            }
-
-        server.show_message_log(f"Optimization completed for {params.functionName} with {speedup:.2f}x speedup", "Info")
-
-        return {
-            "functionName": params.functionName,
-            "status": "success",
-            "message": "Optimization completed successfully",
-            "extra": f"Speedup: {speedup:.2f}x faster",
-            "patch_file": str(patch_path),
-            "task_id": params.task_id,
-            "explanation": best_optimization.explanation_v2,
-        }
+        result = await loop.run_in_executor(None, sync_perform_optimization, server, params)
+    except asyncio.CancelledError:
+        return {"status": "canceled", "message": "Task was canceled"}
+    else:
+        return result
     finally:
         server.cleanup_the_optimizer()