Merge branch 'main' into feat-staging

Author: HeshamHM28

codeflash/code_utils/shell_utils.py

@@ -15,7 +15,7 @@
     SHELL_RC_EXPORT_PATTERN = re.compile(r"^set CODEFLASH_API_KEY=(cf-.*)$", re.MULTILINE)
     SHELL_RC_EXPORT_PREFIX = "set CODEFLASH_API_KEY="
 else:
-    SHELL_RC_EXPORT_PATTERN = re.compile(r'^export CODEFLASH_API_KEY="?(cf-[^\s"]+)"?$', re.MULTILINE)
+    SHELL_RC_EXPORT_PATTERN = re.compile(r'^(?!#)export CODEFLASH_API_KEY=[\'"]?(cf-[^\s"]+)[\'"]$', re.MULTILINE)
     SHELL_RC_EXPORT_PREFIX = "export CODEFLASH_API_KEY="
 
 

codeflash/lsp/beta.py

@@ -36,27 +36,71 @@ def get_optimizable_functions(
     server: CodeflashLanguageServer, params: OptimizableFunctionsParams
 ) -> dict[str, list[str]]:
     file_path = Path(uris.to_fs_path(params.textDocument.uri))
-    server.optimizer.args.file = file_path
-    server.optimizer.args.previous_checkpoint_functions = False
-    optimizable_funcs, _ = server.optimizer.get_optimizable_functions()
-    path_to_qualified_names = {}
-    for path, functions in optimizable_funcs.items():
-        path_to_qualified_names[path.as_posix()] = [func.qualified_name for func in functions]
-    return path_to_qualified_names
+    server.show_message_log(f"Getting optimizable functions for: {file_path}", "Info")
+
+    # Save original args to restore later
+    original_file = getattr(server.optimizer.args, "file", None)
+    original_function = getattr(server.optimizer.args, "function", None)
+    original_checkpoint = getattr(server.optimizer.args, "previous_checkpoint_functions", None)
+
+    server.show_message_log(f"Original args - file: {original_file}, function: {original_function}", "Info")
+
+    try:
+        # Set temporary args for this request only
+        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
+
+        server.show_message_log("Calling get_optimizable_functions...", "Info")
+        optimizable_funcs, _ = server.optimizer.get_optimizable_functions()
+
+        path_to_qualified_names = {}
+        for path, functions in optimizable_funcs.items():
+            path_to_qualified_names[path.as_posix()] = [func.qualified_name for func in functions]
+
+        server.show_message_log(
+            f"Found {len(path_to_qualified_names)} files with functions: {path_to_qualified_names}", "Info"
+        )
+        return path_to_qualified_names
+    finally:
+        # Restore original args to prevent state corruption
+        if original_file is not None:
+            server.optimizer.args.file = original_file
+        if original_function is not None:
+            server.optimizer.args.function = original_function
+        else:
+            server.optimizer.args.function = None
+        if original_checkpoint is not None:
+            server.optimizer.args.previous_checkpoint_functions = original_checkpoint
+
+        server.show_message_log(
+            f"Restored args - file: {server.optimizer.args.file}, function: {server.optimizer.args.function}", "Info"
+        )
 
 
 @server.feature("initializeFunctionOptimization")
 def initialize_function_optimization(
     server: CodeflashLanguageServer, params: FunctionOptimizationParams
 ) -> 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")
+
+    # IMPORTANT: Store the specific function for optimization, but don't corrupt global state
     server.optimizer.args.function = params.functionName
     server.optimizer.args.file = file_path
+
+    server.show_message_log(
+        f"Args set - function: {server.optimizer.args.function}, file: {server.optimizer.args.file}", "Info"
+    )
+
     optimizable_funcs, _ = server.optimizer.get_optimizable_functions()
     if not optimizable_funcs:
+        server.show_message_log(f"No optimizable functions found for {params.functionName}", "Warning")
         return {"functionName": params.functionName, "status": "not found", "args": None}
+
     fto = optimizable_funcs.popitem()[1][0]
     server.optimizer.current_function_being_optimized = fto
+    server.show_message_log(f"Successfully initialized optimization for {params.functionName}", "Info")
     return {"functionName": params.functionName, "status": "success"}
 
 
@@ -136,11 +180,20 @@ def generate_tests(server: CodeflashLanguageServer, params: FunctionOptimization
 
 
 @server.feature("performFunctionOptimization")
-def perform_function_optimization(
+def perform_function_optimization(  # noqa: PLR0911
     server: CodeflashLanguageServer, params: FunctionOptimizationParams
 ) -> dict[str, str]:
+    server.show_message_log(f"Starting optimization for function: {params.functionName}", "Info")
     current_function = server.optimizer.current_function_being_optimized
 
+    if not current_function:
+        server.show_message_log(f"No current function being optimized for {params.functionName}", "Error")
+        return {
+            "functionName": params.functionName,
+            "status": "error",
+            "message": "No function currently being optimized",
+        }
+
     module_prep_result = server.optimizer.prepare_module_for_optimization(current_function.file_path)
 
     validated_original_code, original_module_ast = module_prep_result
@@ -214,25 +267,28 @@ def perform_function_optimization(
     )
 
     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}",
         }
 
     optimized_source = best_optimization.candidate.source_code
+    speedup = original_code_baseline.runtime / best_optimization.runtime
+
+    server.show_message_log(f"Optimization completed for {params.functionName} with {speedup:.2f}x speedup", "Info")
+
+    # CRITICAL: Clear the function filter after optimization to prevent state corruption
+    server.optimizer.args.function = None
+    server.show_message_log("Cleared function filter to prevent state corruption", "Info")
 
     return {
         "functionName": params.functionName,
         "status": "success",
         "message": "Optimization completed successfully",
-        "extra": f"Speedup: {original_code_baseline.runtime / best_optimization.runtime:.2f}x faster",
+        "extra": f"Speedup: {speedup:.2f}x faster",
         "optimization": optimized_source,
     }
-
-
-if __name__ == "__main__":
-    from codeflash.cli_cmds.console import console
-
-    console.quiet = True
-    server.start_io()

codeflash/lsp/server.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
-from lsprotocol.types import INITIALIZE
+from lsprotocol.types import INITIALIZE, MessageType, LogMessageParams
 from pygls import uris
 from pygls.protocol import LanguageServerProtocol, lsp_method
 from pygls.server import LanguageServer
@@ -55,3 +55,24 @@ def initialize_optimizer(self, config_file: Path) -> None:
         args.no_pr = True  # LSP server should not create PRs
         args = process_pyproject_config(args)
         self.optimizer = Optimizer(args)
+
+    def show_message_log(self, message: str, message_type: str) -> None:
+        """Send a log message to the client's output channel.
+        
+        Args:
+            message: The message to log
+            message_type: String type - "Info", "Warning", "Error", or "Log"
+        """
+        # Convert string message type to LSP MessageType enum
+        type_mapping = {
+            "Info": MessageType.Info,
+            "Warning": MessageType.Warning, 
+            "Error": MessageType.Error,
+            "Log": MessageType.Log
+        }
+        
+        lsp_message_type = type_mapping.get(message_type, MessageType.Info)
+        
+        # Send log message to client (appears in output channel)
+        log_params = LogMessageParams(type=lsp_message_type, message=message)
+        self.lsp.notify("window/logMessage", log_params)

codeflash/lsp/server_entry.py

@@ -0,0 +1,48 @@
+"""This script is the dedicated entry point for the Codeflash Language Server.
+It initializes the server and redirects its logs to stderr so that the
+VS Code client can display them in the output channel.
+
+This script is run by the VS Code extension and is not intended to be
+executed directly by users.
+"""
+
+import logging
+import sys
+
+from codeflash.lsp.beta import server
+
+
+# Configure logging to stderr for VS Code output channel
+def setup_logging():
+    # Clear any existing handlers to prevent conflicts
+    root_logger = logging.getLogger()
+    root_logger.handlers.clear()
+
+    # Set up stderr handler for VS Code output channel with [LSP-Server] prefix
+    handler = logging.StreamHandler(sys.stderr)
+    handler.setFormatter(logging.Formatter("[LSP-Server] %(asctime)s [%(levelname)s]: %(message)s"))
+
+    # Configure root logger
+    root_logger.addHandler(handler)
+    root_logger.setLevel(logging.INFO)
+
+    # Also configure the pygls logger specifically
+    pygls_logger = logging.getLogger("pygls")
+    pygls_logger.setLevel(logging.INFO)
+
+    return root_logger
+
+
+if __name__ == "__main__":
+    # Set up logging
+    log = setup_logging()
+    log.info("Starting Codeflash Language Server...")
+
+    # Silence the console module to prevent stdout pollution
+    from codeflash.cli_cmds.console import console
+
+    console.quiet = True
+    # console.enable()
+
+    # Start the language server
+    server.start_io()

codeflash/optimization/function_optimizer.py

@@ -1492,11 +1492,9 @@ def submit_test_generation_tasks(
         ]
 
     def cleanup_generated_files(self) -> None:
-        paths_to_cleanup = [self.test_cfg.concolic_test_root_dir]
+        paths_to_cleanup = []
         for test_file in self.test_files:
             paths_to_cleanup.append(test_file.instrumented_behavior_file_path)
             paths_to_cleanup.append(test_file.benchmarking_file_path)
 
         cleanup_paths(paths_to_cleanup)
-        if hasattr(get_run_tmp_file, "tmpdir"):
-            get_run_tmp_file.tmpdir.cleanup()

codeflash/optimization/optimizer.py

@@ -12,7 +12,7 @@
 from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient
 from codeflash.cli_cmds.console import console, logger, progress_bar
 from codeflash.code_utils import env_utils
-from codeflash.code_utils.code_utils import cleanup_paths
+from codeflash.code_utils.code_utils import cleanup_paths, get_run_tmp_file
 from codeflash.code_utils.env_utils import get_pr_number
 from codeflash.either import is_successful
 from codeflash.models.models import ValidCode
@@ -289,31 +289,36 @@ def run(self) -> None:
                         f"{function_to_optimize.qualified_name}"
                     )
                     console.rule()
+                    function_optimizer = None
+                    try:
+                        function_optimizer = self.create_function_optimizer(
+                            function_to_optimize,
+                            function_to_tests=function_to_tests,
+                            function_to_optimize_source_code=validated_original_code[original_module_path].source_code,
+                            function_benchmark_timings=function_benchmark_timings,
+                            total_benchmark_timings=total_benchmark_timings,
+                            original_module_ast=original_module_ast,
+                            original_module_path=original_module_path,
+                        )
 
-                    function_optimizer = self.create_function_optimizer(
-                        function_to_optimize,
-                        function_to_tests=function_to_tests,
-                        function_to_optimize_source_code=validated_original_code[original_module_path].source_code,
-                        function_benchmark_timings=function_benchmark_timings,
-                        total_benchmark_timings=total_benchmark_timings,
-                        original_module_ast=original_module_ast,
-                        original_module_path=original_module_path,
-                    )
-
-                    self.current_function_optimizer = (
-                        function_optimizer  # needed to clean up from the outside of this function
-                    )
-                    best_optimization = function_optimizer.optimize_function()
-                    if self.functions_checkpoint:
-                        self.functions_checkpoint.add_function_to_checkpoint(
-                            function_to_optimize.qualified_name_with_modules_from_root(self.args.project_root)
+                        self.current_function_optimizer = (
+                            function_optimizer  # needed to clean up from the outside of this function
                         )
-                    if is_successful(best_optimization):
-                        optimizations_found += 1
-                    else:
-                        logger.warning(best_optimization.failure())
-                        console.rule()
-                        continue
+                        best_optimization = function_optimizer.optimize_function()
+                        if self.functions_checkpoint:
+                            self.functions_checkpoint.add_function_to_checkpoint(
+                                function_to_optimize.qualified_name_with_modules_from_root(self.args.project_root)
+                            )
+                        if is_successful(best_optimization):
+                            optimizations_found += 1
+                        else:
+                            logger.warning(best_optimization.failure())
+                            console.rule()
+                            continue
+                    finally:
+                        if function_optimizer is not None:
+                            function_optimizer.cleanup_generated_files()
+
             ph("cli-optimize-run-finished", {"optimizations_found": optimizations_found})
             if self.functions_checkpoint:
                 self.functions_checkpoint.cleanup()
@@ -351,6 +356,9 @@ 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
         cleanup_paths([self.test_cfg.concolic_test_root_dir, self.replay_tests_dir])
 
 

codeflash/telemetry/posthog_cf.py

@@ -7,7 +7,7 @@
 
 from codeflash.api.cfapi import get_user_id
 from codeflash.cli_cmds.console import logger
-from codeflash.version import __version__, __version_tuple__
+from codeflash.version import __version__
 
 _posthog = None
 
@@ -36,7 +36,7 @@ def ph(event: str, properties: dict[str, Any] | None = None) -> None:
         return
 
     properties = properties or {}
-    properties.update({"cli_version": __version__, "cli_version_tuple": __version_tuple__})
+    properties.update({"cli_version": __version__})
 
     user_id = get_user_id()
 

codeflash/update_license_version.py

@@ -2,15 +2,15 @@
 from datetime import datetime
 from pathlib import Path
 
-from version import __version_tuple__
+from .version import __version__
 
 
 def main() -> None:
     # Use the version tuple from version.py
-    version = __version_tuple__
+    version = __version__
 
     # Use the major and minor version components from the version tuple
-    major_minor_version = ".".join(map(str, version[:2]))
+    major_minor_version = ".".join(map(str, version.split(".")[:2]))
 
     # Define the pattern to search for and the replacement string for the version
     version_pattern = re.compile(r"(Licensed Work:\s+Codeflash Client version\s+)(0\.\d+)(\.x)")

codeflash/version.py

@@ -1,3 +1,2 @@
 # These version placeholders will be replaced by uv-dynamic-versioning during build.
-__version__ = "0.14.6"
-__version_tuple__ = (0, 14, 6)
+__version__ = "0.14.7"

pyproject.toml

@@ -268,17 +268,17 @@ enable = true
 style = "pep440"
 vcs = "git"
 
-[tool.uv-dynamic-versioning.substitution]
-files = ["codeflash/version.py"]
-
-[tool.uv-dynamic-versioning.files."codeflash/version.py"]
-persistent-substitution = true
-initial-content = """
-  # These version placeholders will be replaced by uv-dynamic-versioning during build.
-  __version__ = "0.0.0"
-  __version_tuple__ = (0, 0, 0)
+[tool.hatch.build.hooks.version]
+path = "codeflash/version.py"
+template = """# These version placeholders will be replaced by uv-dynamic-versioning during build.
+__version__ = "{version}"
 """
 
+
+#[tool.hatch.build.hooks.custom]
+#path = "codeflash/update_license_version.py"
+
+
 [tool.codeflash]
 module-root = "codeflash"
 tests-root = "tests"