fixes for the api key and the lsp gracefull shutdown

Author: mohammedahmed18

codeflash/api/cfapi.py

@@ -40,6 +40,7 @@ def make_cfapi_request(
     payload: dict[str, Any] | None = None,
     extra_headers: dict[str, str] | None = None,
     *,
+    api_key: str | None = None,
     suppress_errors: bool = False,
 ) -> Response:
     """Make an HTTP request using the specified method, URL, headers, and JSON payload.
@@ -51,7 +52,7 @@ def make_cfapi_request(
     :return: The response object from the API.
     """
     url = f"{CFAPI_BASE_URL}/cfapi{endpoint}"
-    cfapi_headers = {"Authorization": f"Bearer {get_codeflash_api_key()}"}
+    cfapi_headers = {"Authorization": f"Bearer {api_key or get_codeflash_api_key()}"}
     if extra_headers:
         cfapi_headers.update(extra_headers)
     try:
@@ -83,15 +84,17 @@ def make_cfapi_request(
 
 
 @lru_cache(maxsize=1)
-def get_user_id() -> Optional[str]:
+def get_user_id(api_key: Optional[str] = None) -> Optional[str]:
     """Retrieve the user's userid by making a request to the /cfapi/cli-get-user endpoint.
 
     :return: The userid or None if the request fails.
     """
     if not ensure_codeflash_api_key():
         return None
 
-    response = make_cfapi_request(endpoint="/cli-get-user", method="GET", extra_headers={"cli_version": __version__})
+    response = make_cfapi_request(
+        endpoint="/cli-get-user", method="GET", extra_headers={"cli_version": __version__}, api_key=api_key
+    )
     if response.status_code == 200:
         if "min_version" not in response.text:
             return response.text

codeflash/lsp/beta.py

@@ -4,7 +4,7 @@
 import os
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 import git
 from pygls import uris
@@ -104,6 +104,8 @@ def get_optimizable_functions(
 ) -> dict[str, list[str]]:
     file_path = Path(uris.to_fs_path(params.textDocument.uri))
     server.show_message_log(f"Getting optimizable functions for: {file_path}", "Info")
+    if not server.optimizer:
+        return {"status": "error", "message": "optimizer not initialized"}
 
     server.optimizer.args.file = file_path
     server.optimizer.args.function = None  # Always get ALL functions, not just one
@@ -184,8 +186,10 @@ def validate_project(server: CodeflashLanguageServer, _params: FunctionOptimizat
     return {"status": "success", "moduleRoot": args.module_root}
 
 
-def _initialize_optimizer_if_api_key_is_valid(server: CodeflashLanguageServer) -> dict[str, str]:
-    user_id = get_user_id()
+def _initialize_optimizer_if_api_key_is_valid(
+    server: CodeflashLanguageServer, api_key: Optional[str] = None
+) -> dict[str, str]:
+    user_id = get_user_id(api_key=api_key)
     if user_id is None:
         return {"status": "error", "message": "api key not found or invalid"}
 
@@ -224,19 +228,19 @@ def provide_api_key(server: CodeflashLanguageServer, params: ProvideApiKeyParams
         if not api_key.startswith("cf-"):
             return {"status": "error", "message": "Api key is not valid"}
 
-        result = save_api_key_to_rc(api_key)
-        if not is_successful(result):
-            return {"status": "error", "message": result.failure()}
-
         # clear cache to ensure the new api key is used
         get_codeflash_api_key.cache_clear()
         get_user_id.cache_clear()
 
-        init_result = _initialize_optimizer_if_api_key_is_valid(server)
+        init_result = _initialize_optimizer_if_api_key_is_valid(server, api_key)
         if init_result["status"] == "error":
             return {"status": "error", "message": "Api key is not valid"}
 
-        return {"status": "success", "message": "Api key saved successfully", "user_id": init_result["user_id"]}
+        user_id = init_result["user_id"]
+        result = save_api_key_to_rc(api_key)
+        if not is_successful(result):
+            return {"status": "error", "message": result.failure()}
+        return {"status": "success", "message": "Api key saved successfully", "user_id": user_id}  # noqa: TRY300
     except Exception:
         return {"status": "error", "message": "something went wrong while saving the api key"}
 

codeflash/lsp/lsp_logger.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
+import logging
+import sys
 from dataclasses import dataclass
 from typing import Any, Callable, Optional
 
 from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.lsp.lsp_message import LspTextMessage
 
+root_logger = None
+
 
 @dataclass
 class LspMessageTags:
@@ -109,3 +113,27 @@ def enhanced_log(
         clean_msg = LspTextMessage(text=clean_msg, takes_time=final_tags.loading).serialize()
 
     actual_log_fn(clean_msg, *args, **kwargs)
+
+
+# Configure logging to stderr for VS Code output channel
+def setup_logging() -> logging.Logger:
+    global root_logger  # noqa: PLW0603
+    if root_logger:
+        return root_logger
+    # Clear any existing handlers to prevent conflicts
+    logger = logging.getLogger()
+    logger.handlers.clear()
+
+    # Set up stderr handler for VS Code output channel with [LSP-Server] prefix
+    handler = logging.StreamHandler(sys.stderr)
+    handler.setLevel(logging.DEBUG)
+
+    # Configure root logger
+    logger.addHandler(handler)
+
+    # Also configure the pygls logger specifically
+    pygls_logger = logging.getLogger("pygls")
+    pygls_logger.setLevel(logging.INFO)
+
+    root_logger = logger
+    return logger

codeflash/lsp/server.py

@@ -1,14 +1,12 @@
 from __future__ import annotations
 
-import sys
 from pathlib import Path
-from threading import Event
-from typing import TYPE_CHECKING, Any, Optional, TextIO
+from typing import TYPE_CHECKING, Any
 
 from lsprotocol.types import INITIALIZE, LogMessageParams, MessageType
 from pygls import uris
 from pygls.protocol import LanguageServerProtocol, lsp_method
-from pygls.server import LanguageServer, StdOutTransportAdapter, aio_readline
+from pygls.server import LanguageServer
 
 if TYPE_CHECKING:
     from lsprotocol.types import InitializeParams, InitializeResult
@@ -85,6 +83,8 @@ def show_message_log(self, message: str, message_type: str) -> None:
         self.lsp.notify("window/logMessage", log_params)
 
     def cleanup_the_optimizer(self) -> None:
+        if not self.optimizer:
+            return
         try:
             self.optimizer.cleanup_temporary_paths()
             # restore args and test cfg
@@ -96,26 +96,7 @@ def cleanup_the_optimizer(self) -> None:
         except Exception:
             self.show_message_log("Failed to cleanup optimizer", "Error")
 
-    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()
+    def shutdown(self) -> None:
+        """Gracefully shutdown the server."""
+        self.cleanup_the_optimizer()
+        super().shutdown()

codeflash/lsp/server_entry.py

@@ -7,36 +7,13 @@
 executed directly by users.
 """
 
-import logging
-import sys
-
 from codeflash.lsp.beta import server
-
-
-# Configure logging to stderr for VS Code output channel
-def setup_logging() -> logging.Logger:
-    # Clear any existing handlers to prevent conflicts
-    root_logger = logging.getLogger()
-    root_logger.handlers.clear()
-
-    # Set up stderr handler for VS Code output channel with [LSP-Server] prefix
-    handler = logging.StreamHandler(sys.stderr)
-    handler.setLevel(logging.DEBUG)
-
-    # Configure root logger
-    root_logger.addHandler(handler)
-
-    # Also configure the pygls logger specifically
-    pygls_logger = logging.getLogger("pygls")
-    pygls_logger.setLevel(logging.INFO)
-
-    return root_logger
-
+from codeflash.lsp.lsp_logger import setup_logging
 
 if __name__ == "__main__":
     # Set up logging
-    log = setup_logging()
-    log.info("Starting Codeflash Language Server...")
+    root_logger = setup_logging()
+    root_logger.info("Starting Codeflash Language Server...")
 
     # Start the language server
     server.start_io()

codeflash/optimization/function_optimizer.py

@@ -164,7 +164,8 @@ def _process_line_profiler_results(self) -> OptimizedCandidate | None:
 
     def _process_refinement_results(self) -> OptimizedCandidate | None:
         """Process refinement results and add to queue."""
-        logger.info("loading|Refining generated code for improved quality and performance...")
+        if self.future_all_refinements:
+            logger.info("loading|Refining generated code for improved quality and performance...")
         concurrent.futures.wait(self.future_all_refinements)
         refinement_response = []