move optimizer cleanup into server and run optimization in thread

Author: mohammedahmed18

codeflash/lsp/beta.py

@@ -110,7 +110,7 @@ def initialize_function_optimization(
 
     if count == 0:
         server.show_message_log(f"No optimizable functions found for {params.functionName}", "Warning")
-        cleanup_the_optimizer(server)
+        server.cleanup_the_optimizer()
         return {"functionName": params.functionName, "status": "error", "message": "not found", "args": None}
 
     fto = optimizable_funcs.popitem()[1][0]
@@ -217,6 +217,7 @@ def provide_api_key(server: CodeflashLanguageServer, params: ProvideApiKeyParams
 
 
 @server.feature("performFunctionOptimization")
+@server.thread()
 def perform_function_optimization(  # noqa: PLR0911
     server: CodeflashLanguageServer, params: FunctionOptimizationParams
 ) -> dict[str, str]:
@@ -337,14 +338,4 @@ def perform_function_optimization(  # noqa: PLR0911
             "explanation": best_optimization.explanation_v2,
         }
     finally:
-        cleanup_the_optimizer(server)
-
-
-def cleanup_the_optimizer(server: CodeflashLanguageServer) -> None:
-    server.optimizer.cleanup_temporary_paths()
-    # restore args and test cfg
-    if server.optimizer.original_args_and_test_cfg:
-        server.optimizer.args, server.optimizer.test_cfg = server.optimizer.original_args_and_test_cfg
-    server.optimizer.args.function = None
-    server.optimizer.current_worktree = None
-    server.optimizer.current_function_optimizer = None
+        server.cleanup_the_optimizer()

codeflash/lsp/server.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+import sys
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from threading import Event
+from typing import TYPE_CHECKING, Any, Optional, TextIO
 
 from lsprotocol.types import INITIALIZE, LogMessageParams, MessageType
 from pygls import uris
 from pygls.protocol import LanguageServerProtocol, lsp_method
-from pygls.server import LanguageServer
+from pygls.server import LanguageServer, StdOutTransportAdapter, aio_readline
 
 if TYPE_CHECKING:
     from lsprotocol.types import InitializeParams, InitializeResult
@@ -81,3 +83,36 @@ def show_message_log(self, message: str, message_type: str) -> None:
         # 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)
+
+    def cleanup_the_optimizer(self) -> None:
+        self.optimizer.cleanup_temporary_paths()
+        # restore args and test cfg
+        if self.optimizer.original_args_and_test_cfg:
+            self.optimizer.args, self.optimizer.test_cfg = self.optimizer.original_args_and_test_cfg
+        self.optimizer.args.function = None
+        self.optimizer.current_worktree = None
+        self.optimizer.current_function_optimizer = None
+
+    def start_io(self, stdin: Optional[TextIO] = None, stdout: Optional[TextIO] = None) -> None:
+        self.show_message_log("Starting IO server", "Info")
+
+        self._stop_event = Event()
+        transport = StdOutTransportAdapter(stdin or sys.stdin.buffer, stdout or sys.stdout.buffer)
+        self.lsp.connection_made(transport)
+        try:
+            self.loop.run_until_complete(
+                aio_readline(
+                    self.loop,
+                    self.thread_pool_executor,
+                    self._stop_event,
+                    stdin or sys.stdin.buffer,
+                    self.lsp.data_received,
+                )
+            )
+        except BrokenPipeError:
+            self.show_message_log("Connection to the client is lost! Shutting down the server.", "Error")
+        except (KeyboardInterrupt, SystemExit):
+            pass
+        finally:
+            self.cleanup_the_optimizer()
+            self.shutdown()