second step

Author: KRRT7

codeflash/code_utils/code_utils.py

@@ -216,7 +216,5 @@ def restore_conftest(path_to_content_map: dict[Path, str]) -> None:
         path.write_text(file_content, encoding="utf8")
 
 
-def exit_with_message(message: str, *, error_on_exit: bool = False) -> None:
-    paneled_text(message, panel_args={"style": "red"})
-
-    sys.exit(1 if error_on_exit else 0)
+async def dummy_async_function() -> None:
+    """Provide a dummy async function for testing purposes."""

codeflash/lsp/beta.py

@@ -6,6 +6,7 @@
 
 from pygls import uris
 
+from codeflash.either import is_successful
 from codeflash.lsp.server import CodeflashLanguageServer, CodeflashLanguageServerProtocol
 
 if TYPE_CHECKING:
@@ -80,30 +81,23 @@ def third_step_in_optimize_function(server: CodeflashLanguageServer, params: Opt
     )
 
     server.optimizer.current_function_optimizer = function_optimizer
+    if not function_optimizer:
+        return {"functionName": params.functionName, "status": "error", "message": "No function optimizer found"}
+
+    initialization_result = function_optimizer.can_be_optimized()
+    if not is_successful(initialization_result):
+        return {"functionName": params.functionName, "status": "error", "message": initialization_result.failure()}
+
+    should_run_experiment, code_context, original_helper_code = initialization_result.unwrap()
 
     return {
         "functionName": params.functionName,
         "status": "success",
-        "message": "Function optimizer created successfully",
-        "extra": function_optimizer.function_to_tests,
+        "message": "Function can be optimized",
+        "extra": original_helper_code,
     }
 
 
-@server.feature("fourth_step_in_optimize_function")
-def fourth_step_in_optimize_function(server: CodeflashLanguageServer, params: OptimizeFunctionParams) -> dict[str, str]:
-    current_function_optimizer = server.optimizer.current_function_optimizer
-
-    if not current_function_optimizer:
-        return {"functionName": params.functionName, "status": "error", "message": "No function optimizer found"}
-
-    optimized_code = current_function_optimizer.optimize_function()
-
-    if not optimized_code:
-        return {"functionName": params.functionName, "status": "error", "message": "Optimization failed"}
-
-    return {"functionName": params.functionName, "status": "success", "optimized_code": optimized_code}
-
-
 if __name__ == "__main__":
     from codeflash.cli_cmds.console import console
 

codeflash/optimization/function_optimizer.py

@@ -144,7 +144,7 @@ def __init__(
         self.total_benchmark_timings = total_benchmark_timings if total_benchmark_timings else {}
         self.replay_tests_dir = replay_tests_dir if replay_tests_dir else None
 
-    def optimize_function(self) -> Result[BestOptimization, str]:  # noqa: PLR0911
+    def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[Path, str]], str]:
         should_run_experiment = self.experiment_id is not None
         logger.debug(f"Function Trace ID: {self.function_trace_id}")
         ph("cli-optimize-function-start", {"function_trace_id": self.function_trace_id})
@@ -171,6 +171,15 @@ def optimize_function(self) -> Result[BestOptimization, str]:  # noqa: PLR0911
         ):
             return Failure("Function optimization previously attempted, skipping.")
 
+        return Success((should_run_experiment, code_context, original_helper_code))
+
+    def optimize_function(self) -> Result[BestOptimization, str]:
+        initialization_result = self.can_be_optimized()
+        if not is_successful(initialization_result):
+            return Failure(initialization_result.failure())
+
+        should_run_experiment, code_context, original_helper_code = initialization_result.unwrap()
+
         code_print(code_context.read_writable_code)
         generated_test_paths = [
             get_test_file_path(