lsp: verify and provide api key

Author: mohammedahmed18

codeflash/api/cfapi.py

@@ -99,6 +99,9 @@ def get_user_id() -> Optional[str]:
             if min_version and version.parse(min_version) > version.parse(__version__):
                 msg = "Your Codeflash CLI version is outdated. Please update to the latest version using `pip install --upgrade codeflash`."
                 console.print(f"[bold red]{msg}[/bold red]")
+                if console.quiet:  # lsp
+                    logger.debug(msg)
+                    return None
                 sys.exit(1)
             return userid
 

codeflash/code_utils/env_utils.py

@@ -34,7 +34,12 @@ def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool =
 
 @lru_cache(maxsize=1)
 def get_codeflash_api_key() -> str:
-    api_key = os.environ.get("CODEFLASH_API_KEY") or read_api_key_from_shell_config()
+    if console.quiet:  # lsp
+        # prefer shell config over env var in lsp mode
+        api_key = read_api_key_from_shell_config() or os.environ.get("CODEFLASH_API_KEY")
+    else:
+        api_key = os.environ.get("CODEFLASH_API_KEY") or read_api_key_from_shell_config()
+
     api_secret_docs_message = "For more information, refer to the documentation at [https://docs.codeflash.ai/getting-started/codeflash-github-actions#add-your-api-key-to-your-repository-secrets]."  # noqa
     if not api_key:
         msg = (

codeflash/lsp/beta.py

@@ -8,6 +8,8 @@
 
 from pygls import uris
 
+from codeflash.api.cfapi import 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
 
@@ -28,6 +30,11 @@ class FunctionOptimizationParams:
     functionName: str  # noqa: N815
 
 
+@dataclass
+class ProvideApiKeyParams:
+    api_key: str
+
+
 server = CodeflashLanguageServer("codeflash-language-server", "v1.0", protocol_cls=CodeflashLanguageServerProtocol)
 
 
@@ -118,6 +125,40 @@ 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"}
+
+        from codeflash.optimization.optimizer import Optimizer
+
+        server.optimizer = Optimizer(server.args)
+        return {"status": "success", "user_id": user_id}  # noqa
+    except Exception:
+        return {"status": "error", "message": "something went wrong while validating the api key"}
+
+
+@server.feature("provideApiKey")
+def provide_api_key(server: CodeflashLanguageServer, params: ProvideApiKeyParams) -> dict[str, str]:
+    try:
+        api_key = params.api_key
+        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
+
+        server.optimizer = Optimizer(server.args)
+        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"}
+
+
 @server.feature("prepareOptimization")
 def prepare_optimization(server: CodeflashLanguageServer, params: FunctionOptimizationParams) -> dict[str, str]:
     current_function = server.optimizer.current_function_being_optimized

codeflash/lsp/server.py

@@ -25,7 +25,7 @@ def lsp_initialize(self, params: InitializeParams) -> InitializeResult:
             workspace_path = uris.to_fs_path(workspace_uri)
             pyproject_toml_path = self._find_pyproject_toml(workspace_path)
             if pyproject_toml_path:
-                server.initialize_optimizer(pyproject_toml_path)
+                server.prepare_optimizer_arguments(pyproject_toml_path)
                 server.show_message(f"Found pyproject.toml at: {pyproject_toml_path}")
             else:
                 server.show_message("No pyproject.toml found in workspace.")
@@ -45,16 +45,17 @@ class CodeflashLanguageServer(LanguageServer):
     def __init__(self, *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
         super().__init__(*args, **kwargs)
         self.optimizer = None
+        self.args = None
 
-    def initialize_optimizer(self, config_file: Path) -> None:
+    def prepare_optimizer_arguments(self, config_file: Path) -> None:
         from codeflash.cli_cmds.cli import parse_args, process_pyproject_config
-        from codeflash.optimization.optimizer import Optimizer
 
         args = parse_args()
         args.config_file = config_file
         args.no_pr = True  # LSP server should not create PRs
         args = process_pyproject_config(args)
-        self.optimizer = Optimizer(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
 
     def show_message_log(self, message: str, message_type: str) -> None:
         """Send a log message to the client's output channel.