diff --git a/codeflash/lsp/beta.py b/codeflash/lsp/beta.py
index c4e305faa..14d43a1a5 100644
--- a/codeflash/lsp/beta.py
+++ b/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()
 
diff --git a/codeflash/lsp/features/perform_optimization.py b/codeflash/lsp/features/perform_optimization.py
index 9825c5802..5dfd0daef 100644
--- a/codeflash/lsp/features/perform_optimization.py
+++ b/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,
diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py
index a0dedb570..2eef51f0f 100644
--- a/codeflash/optimization/function_optimizer.py
+++ b/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",