fix: clear cache before checking the api key to pick the latest one, and enable codeflash debug messages in lsp

Author: mohammedahmed18

codeflash/code_utils/shell_utils.py

@@ -42,7 +42,7 @@ def get_shell_rc_path() -> Path:
 
 
 def get_api_key_export_line(api_key: str) -> str:
-    return f"{SHELL_RC_EXPORT_PREFIX}{api_key}"
+    return f'{SHELL_RC_EXPORT_PREFIX}"{api_key}"'
 
 
 def save_api_key_to_rc(api_key: str) -> Result[str, str]:

codeflash/lsp/beta.py

@@ -8,7 +8,7 @@
 
 from pygls import uris
 
-from codeflash.api.cfapi import get_user_id
+from codeflash.api.cfapi import get_codeflash_api_key, get_user_id
 from codeflash.code_utils.shell_utils import save_api_key_to_rc
 from codeflash.either import is_successful
 from codeflash.lsp.server import CodeflashLanguageServer, CodeflashLanguageServerProtocol
@@ -125,17 +125,21 @@ def discover_function_tests(server: CodeflashLanguageServer, params: FunctionOpt
     return {"functionName": params.functionName, "status": "success", "discovered_tests": num_discovered_tests}
 
 
-@server.feature("apiKeyExistsAndValid")
-def check_api_key(_server: CodeflashLanguageServer, _params: any) -> dict[str, str]:
-    try:
-        user_id = get_user_id()
-        if user_id is None:
-            return {"status": "error", "message": "api key not found or invalid"}
+def _initialize_optimizer_if_valid(server: CodeflashLanguageServer) -> dict[str, str]:
+    user_id = get_user_id()
+    if user_id is None:
+        return {"status": "error", "message": "api key not found or invalid"}
+
+    from codeflash.optimization.optimizer import Optimizer
+
+    server.optimizer = Optimizer(server.args)
+    return {"status": "success", "user_id": user_id}
 
-        from codeflash.optimization.optimizer import Optimizer
 
-        server.optimizer = Optimizer(server.args)
-        return {"status": "success", "user_id": user_id}  # noqa
+@server.feature("apiKeyExistsAndValid")
+def check_api_key(server: CodeflashLanguageServer, _params: any) -> dict[str, str]:
+    try:
+        return _initialize_optimizer_if_valid(server)
     except Exception:
         return {"status": "error", "message": "something went wrong while validating the api key"}
 
@@ -144,17 +148,22 @@ def check_api_key(_server: CodeflashLanguageServer, _params: any) -> dict[str, s
 def provide_api_key(server: CodeflashLanguageServer, params: ProvideApiKeyParams) -> dict[str, str]:
     try:
         api_key = params.api_key
+        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()}
-        user_id = get_user_id()
-        if user_id is None:
-            return {"status": "error", "message": "api key is not valid"}
 
-        from codeflash.optimization.optimizer import Optimizer
+        # 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_valid(server)
+        if init_result["status"] == "error":
+            return {"status": "error", "message": "Api key is not valid"}
 
-        server.optimizer = Optimizer(server.args)
-        return {"status": "success", "message": "Api key saved successfully", "user_id": user_id}  # noqa: TRY300
+        return {"status": "success", "message": "Api key saved successfully", "user_id": init_result["user_id"]}
     except Exception:
         return {"status": "error", "message": "something went wrong while saving the api key"}
 

codeflash/lsp/server.py

@@ -55,7 +55,7 @@ def prepare_optimizer_arguments(self, config_file: Path) -> None:
         args.no_pr = True  # LSP server should not create PRs
         args = process_pyproject_config(args)
         self.args = args
-        # self.optimizer = Optimizer(args) # avoid creating the optimizer during initialization, bacause it may cause and erro if the api key is invalid
+        # avoid initializing the optimizer during initialization, because it can cause an error if the api key is invalid
 
     def show_message_log(self, message: str, message_type: str) -> None:
         """Send a log message to the client's output channel.
@@ -71,6 +71,7 @@ def show_message_log(self, message: str, message_type: str) -> None:
             "Warning": MessageType.Warning,
             "Error": MessageType.Error,
             "Log": MessageType.Log,
+            "Debug": MessageType.Debug,
         }
 
         lsp_message_type = type_mapping.get(message_type, MessageType.Info)