Merge branch 'main' into trace-and-optimize

KRRT7 · web-flow · commit 09bf156651e4 · 2025-06-27T21:16:58.000Z
diff --git a/.github/workflows/e2e-bubblesort-benchmark.yaml b/.github/workflows/e2e-bubblesort-benchmark.yaml
@@ -1,4 +1,4 @@
-name: end-to-end-test
+name: E2E - Bubble Sort Benchmark
 
 on:
   pull_request:
diff --git a/.github/workflows/e2e-bubblesort-pytest-nogit.yaml b/.github/workflows/e2e-bubblesort-pytest-nogit.yaml
@@ -1,4 +1,4 @@
-name: end-to-end-test
+name: E2E - Bubble Sort Pytest (No Git)
 
 on:
   pull_request:
diff --git a/.github/workflows/e2e-bubblesort-unittest.yaml b/.github/workflows/e2e-bubblesort-unittest.yaml
@@ -1,4 +1,4 @@
-name: end-to-end-test
+name: E2E - Bubble Sort Unittest
 
 on:
   pull_request:
diff --git a/.github/workflows/e2e-coverage-optimization.yaml b/.github/workflows/e2e-coverage-optimization.yaml
diff --git a/.github/workflows/e2e-futurehouse-structure.yaml b/.github/workflows/e2e-futurehouse-structure.yaml
@@ -1,4 +1,4 @@
-name: end-to-end-test
+name: E2E - Futurehouse Structure
 
 on:
   pull_request:
diff --git a/.github/workflows/e2e-init-optimization.yaml b/.github/workflows/e2e-init-optimization.yaml
@@ -1,4 +1,4 @@
-name: end-to-end-test
+name: E2E - Init Optimization
 
 on:
   pull_request:
diff --git a/.github/workflows/e2e-topological-sort.yaml b/.github/workflows/e2e-topological-sort.yaml
@@ -1,4 +1,4 @@
-name: end-to-end-test
+name: E2E - Topological Sort
 
 on:
   pull_request:
diff --git a/.github/workflows/e2e-tracer-replay.yaml b/.github/workflows/e2e-tracer-replay.yaml
@@ -1,4 +1,4 @@
-name: end-to-end-test
+name: E2E - Tracer Replay
 
 on:
   pull_request:
diff --git a/codeflash.code-workspace b/codeflash.code-workspace
@@ -10,13 +10,16 @@
         }
     ],
     "settings": {
-        "python.defaultInterpreterPath": "~/miniforge3/envs/codeflash312/bin/python",
         "python.terminal.activateEnvironment": true,
         "python.testing.pytestEnabled": true,
-        "python.testing.pytestArgs": ["tests/", "-vv"],
+        "python.testing.pytestArgs": [
+            "tests/",
+            "-vv",
+            "--ignore",
+            "tests/benchmarks/"
+        ],
     },
     "launch": {
-        "version": "0.2.0",
         "configurations": [
             {
                 "name": "bubble_sort",
@@ -62,19 +65,26 @@
                 },
             },
             {
-                "name": "bubble_sort --file bubble_sort.py (MBR)",
+                "name": "run codeflash in custom directory",
                 "type": "debugpy",
                 "request": "launch",
                 "program": "${workspaceFolder:codeflash}/codeflash/main.py",
                 "args": [
                     "--all",
                 ],
-                "cwd": "/Users/krrt7/Desktop/work/my-best-repo",
+                "cwd": "${input:chooseCwd}",
                 "console": "integratedTerminal",
                 "env": {
                     "PYTHONUNBUFFERED": "1"
-                },
+                }
+            }
+        ],
+        "inputs": [
+            {
+                "id": "chooseCwd",
+                "type": "promptString",
+                "description": "Enter the working directory for codeflash",
             }
         ]
-    }
+    },
 }
diff --git a/codeflash/api/aiservice.py b/codeflash/api/aiservice.py
@@ -10,7 +10,7 @@
 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
+from codeflash.code_utils.env_utils import get_codeflash_api_key, is_LSP_enabled
 from codeflash.code_utils.git_utils import get_last_commit_author_if_pr_exists, get_repo_owner_and_name
 from codeflash.models.models import OptimizedCandidate
 from codeflash.telemetry.posthog_cf import ph
@@ -182,6 +182,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
             "python_version": platform.python_version(),
             "experiment_metadata": experiment_metadata,
             "codeflash_version": codeflash_version,
+            "lsp_mode": is_LSP_enabled(),
         }
 
         logger.info("Generating optimized candidates…")
diff --git a/codeflash/code_utils/env_utils.py b/codeflash/code_utils/env_utils.py
@@ -7,7 +7,7 @@
 from pathlib import Path
 from typing import Any, Optional
 
-from codeflash.cli_cmds.console import logger
+from codeflash.cli_cmds.console import console, 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
@@ -110,3 +110,8 @@ def get_cached_gh_event_data() -> dict[str, Any] | None:
         return None
     with Path(event_path).open() as f:
         return json.load(f)  # type: ignore  # noqa
+
+
+@lru_cache(maxsize=1)
+def is_LSP_enabled() -> bool:
+    return console.quiet
diff --git a/codeflash/code_utils/formatter.py b/codeflash/code_utils/formatter.py
@@ -109,6 +109,9 @@ def format_code(
     print_status: bool = True,  # noqa
     exit_on_failure: bool = True,  # noqa
 ) -> str:
+    if console.quiet:
+        # lsp mode
+        exit_on_failure = False
     with tempfile.TemporaryDirectory() as test_dir_str:
         if isinstance(path, str):
             path = Path(path)
diff --git a/codeflash/lsp/beta.py b/codeflash/lsp/beta.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import contextlib
+import os
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -55,18 +57,21 @@ def initialize_function_optimization(
         return {"functionName": params.functionName, "status": "not found", "args": None}
     fto = optimizable_funcs.popitem()[1][0]
     server.optimizer.current_function_being_optimized = fto
-    return {"functionName": params.functionName, "status": "success", "info": fto.server_info}
+    return {"functionName": params.functionName, "status": "success"}
 
 
 @server.feature("discoverFunctionTests")
 def discover_function_tests(server: CodeflashLanguageServer, params: FunctionOptimizationParams) -> dict[str, str]:
-    current_function = server.optimizer.current_function_being_optimized
+    fto = server.optimizer.current_function_being_optimized
+    optimizable_funcs = {fto.file_path: [fto]}
+
+    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)
 
-    optimizable_funcs = {current_function.file_path: [current_function]}
+    server.optimizer.discovered_tests = function_to_tests
 
-    function_to_tests, num_discovered_tests = server.optimizer.discover_tests(optimizable_funcs)
-    # mocking in order to get things going
-    return {"functionName": params.functionName, "status": "success", "generated_tests": str(num_discovered_tests)}
+    return {"functionName": params.functionName, "status": "success", "discovered_tests": num_discovered_tests}
 
 
 @server.feature("prepareOptimization")
@@ -145,6 +150,7 @@ def perform_function_optimization(
         function_to_optimize_source_code=validated_original_code[current_function.file_path].source_code,
         original_module_ast=original_module_ast,
         original_module_path=current_function.file_path,
+        function_to_tests=server.optimizer.discovered_tests or {},
     )
 
     server.optimizer.current_function_optimizer = function_optimizer
@@ -214,13 +220,14 @@ def perform_function_optimization(
             "message": f"No best optimizations found for function {function_to_optimize_qualified_name}",
         }
 
-    optimized_source = best_optimization.candidate.source_code  # noqa: F841
+    optimized_source = best_optimization.candidate.source_code
 
     return {
         "functionName": params.functionName,
         "status": "success",
         "message": "Optimization completed successfully",
         "extra": f"Speedup: {original_code_baseline.runtime / best_optimization.runtime:.2f}x faster",
+        "optimization": optimized_source,
     }
 
 
diff --git a/codeflash/lsp/server.py b/codeflash/lsp/server.py
@@ -52,5 +52,6 @@ def initialize_optimizer(self, config_file: Path) -> None:
 
         args = parse_args()
         args.config_file = config_file
+        args.no_pr = True  # LSP server should not create PRs
         args = process_pyproject_config(args)
         self.optimizer = Optimizer(args)
diff --git a/codeflash/result/critic.py b/codeflash/result/critic.py
@@ -56,8 +56,8 @@ def quantity_of_tests_critic(candidate_result: OptimizedCandidateResult) -> bool
 
     if pass_count >= MIN_TESTCASE_PASSED_THRESHOLD:
         return True
-    # If only one test passed, check if it's a REPLAY_TEST
-    return bool(pass_count == 1 and report[TestType.REPLAY_TEST]["passed"] == 1)
+    # If one or more tests passed, check if least one of them was a successful REPLAY_TEST
+    return bool(pass_count >= 1 and report[TestType.REPLAY_TEST]["passed"] >= 1)
 
 
 def coverage_critic(original_code_coverage: CoverageData | None, test_framework: str) -> bool:
diff --git a/codeflash/verification/concolic_testing.py b/codeflash/verification/concolic_testing.py
@@ -10,6 +10,7 @@
 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.telemetry.posthog_cf import ph
@@ -28,6 +29,11 @@ def generate_concolic_tests(
     start_time = time.perf_counter()
     function_to_concolic_tests = {}
     concolic_test_suite_code = ""
+
+    if is_LSP_enabled():
+        logger.debug("Skipping concolic test generation in LSP mode")
+        return function_to_concolic_tests, concolic_test_suite_code
+
     if (
         test_cfg.concolic_test_root_dir
         and isinstance(function_to_optimize_ast, (ast.FunctionDef, ast.AsyncFunctionDef))
@@ -43,7 +49,7 @@ def generate_concolic_tests(
                     "crosshair",
                     "cover",
                     "--example_output_format=pytest",
-                    "--per_condition_timeout=64",
+                    "--per_condition_timeout=20",
                     ".".join(
                         [
                             function_to_optimize.file_path.relative_to(args.project_root)

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-name: end-to-end-test`
	`1`	`+name: E2E - Bubble Sort Benchmark`
`2`	`2`
`3`	`3`	`on:`
`4`	`4`	`pull_request:`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-name: end-to-end-test`
	`1`	`+name: E2E - Bubble Sort Pytest (No Git)`
`2`	`2`
`3`	`3`	`on:`
`4`	`4`	`pull_request:`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-name: end-to-end-test`
	`1`	`+name: E2E - Bubble Sort Unittest`
`2`	`2`
`3`	`3`	`on:`
`4`	`4`	`pull_request:`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-name: end-to-end-test`
	`1`	`+name: E2E - Futurehouse Structure`
`2`	`2`
`3`	`3`	`on:`
`4`	`4`	`pull_request:`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-name: end-to-end-test`
	`1`	`+name: E2E - Init Optimization`
`2`	`2`
`3`	`3`	`on:`
`4`	`4`	`pull_request:`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-name: end-to-end-test`
	`1`	`+name: E2E - Topological Sort`
`2`	`2`
`3`	`3`	`on:`
`4`	`4`	`pull_request:`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-name: end-to-end-test`
	`1`	`+name: E2E - Tracer Replay`
`2`	`2`
`3`	`3`	`on:`
`4`	`4`	`pull_request:`