Merge pull request #942 from codeflash-ai/lsp/retrun-runtime-info-and-cancel-demo-opt

[LSP] Runtime info for sucessful opt & cancel demo optimization

Author: mohammedahmed18

codeflash/lsp/beta.py

@@ -491,6 +491,7 @@ def initialize_function_optimization(params: FunctionOptimizationInitParams) ->
 async def start_demo_optimization(params: DemoOptimizationParams) -> dict[str, str]:
     try:
         _init()
+        cancel_event = threading.Event()
         # start by creating the worktree so that the demo file is not created in user workspace
         server.optimizer.worktree_mode()
         file_path = create_find_common_tags_file(server.args, params.functionName + ".py")
@@ -509,6 +510,9 @@ async def start_demo_optimization(params: DemoOptimizationParams) -> dict[str, s
         return await perform_function_optimization(
             FunctionOptimizationParams(functionName=params.functionName, task_id=None)
         )
+    except asyncio.CancelledError:
+        cancel_event.set()
+        return get_cancelled_reponse()
     finally:
         server.cleanup_the_optimizer()
 

codeflash/lsp/features/perform_optimization.py

@@ -131,6 +131,8 @@ def sync_perform_optimization(server: CodeflashLanguageServer, cancel_event: thr
         "status": "success",
         "message": "Optimization completed successfully",
         "extra": f"Speedup: {speedup:.2f}x faster",
+        "original_runtime": original_code_baseline.runtime,
+        "optimized_runtime": best_optimization.runtime,
         "patch_file": str(patch_path),
         "task_id": params.task_id,
         "explanation": best_optimization.explanation_v2,

codeflash/optimization/function_optimizer.py

@@ -846,7 +846,7 @@ def log_successful_optimization(
             md_lines = [
                 "### ⚡️ Optimization Summary",
                 f"Function: `{self.function_to_optimize.qualified_name}`",
-                f"File: `{explanation.file_path}`",
+                f"File: {explanation.file_path}",
                 f"Performance: {explanation.perf_improvement_line}",
                 "",
                 "#### Explanation\n",