cleanup

Author: mohammedahmed18

codeflash/lsp/beta.py

@@ -289,11 +289,12 @@ def process_args() -> Namespace:
 
 
 def _init() -> Namespace:
-    if not server.initialized:
-        new_args = process_args()
-        _initialize_optimizer(new_args)
-        server.initialized = True
-    return server.args
+    if server.initialized:
+        return server.args
+    new_args = process_args()
+    _initialize_optimizer(new_args)
+    server.initialized = True
+    return new_args
 
 
 @server.feature("apiKeyExistsAndValid")
@@ -417,13 +418,13 @@ def initialize_function_optimization(params: FunctionOptimizationInitParams) ->
 async def perform_function_optimization(params: FunctionOptimizationParams) -> dict[str, str]:
     with execution_context(task_id=params.task_id):
         loop = asyncio.get_running_loop()
-        server.cancel_event = threading.Event()
+        cancel_event = threading.Event()
 
         try:
             ctx = contextvars.copy_context()
-            return await loop.run_in_executor(None, ctx.run, sync_perform_optimization, server, params)
+            return await loop.run_in_executor(None, ctx.run, sync_perform_optimization, server, cancel_event, params)
         except asyncio.CancelledError:
-            server.cancel_event.set()
+            cancel_event.set()
             return get_cancelled_reponse()
         finally:
             server.cleanup_the_optimizer()

codeflash/lsp/features/perform_optimization.py

@@ -23,8 +23,7 @@ def abort_if_cancelled(cancel_event: threading.Event) -> None:
         raise RuntimeError("cancelled")
 
 
-def sync_perform_optimization(server: CodeflashLanguageServer, params) -> dict[str, str]:  # noqa
-    cancel_event = server.cancel_event
+def sync_perform_optimization(server: CodeflashLanguageServer, cancel_event: threading.Event, params) -> dict[str, str]:  # noqa
     server.show_message_log(f"Starting optimization for function: {params.functionName}", "Info")
     should_run_experiment, code_context, original_helper_code = server.current_optimization_init_result
     function_optimizer = server.optimizer.current_function_optimizer

codeflash/lsp/lsp_message.py

@@ -36,7 +36,8 @@ def type(self) -> str:
     def serialize(self) -> str:
         from codeflash.lsp.beta import server
 
-        current_task_id = server.execution_context_vars.get().get("task_id", None)
+        execution_ctx = server.execution_context_vars.get()
+        current_task_id = execution_ctx.get("task_id", None)
         data = self._loop_through(asdict(self))
         ordered = {"type": self.type(), "task_id": current_task_id, **data}
         return message_delimiter + json.dumps(ordered) + message_delimiter