From 38e4e02148f187dee1625e6811d24616705b7c1f Mon Sep 17 00:00:00 2001 From: mohammed Date: Wed, 9 Jul 2025 15:28:59 +0300 Subject: [PATCH 1/3] lsp: verify and provide api key --- codeflash/api/cfapi.py | 3 +++ codeflash/code_utils/env_utils.py | 7 +++++- codeflash/lsp/beta.py | 41 +++++++++++++++++++++++++++++++ codeflash/lsp/server.py | 9 ++++--- 4 files changed, 55 insertions(+), 5 deletions(-) diff --git a/codeflash/api/cfapi.py b/codeflash/api/cfapi.py index 0f1fbf654..958375d88 100644 --- a/codeflash/api/cfapi.py +++ b/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 diff --git a/codeflash/code_utils/env_utils.py b/codeflash/code_utils/env_utils.py index 2c8ca22a1..476edebdf 100644 --- a/codeflash/code_utils/env_utils.py +++ b/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 = ( diff --git a/codeflash/lsp/beta.py b/codeflash/lsp/beta.py index 2fc91f0a0..10335c550 100644 --- a/codeflash/lsp/beta.py +++ b/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 diff --git a/codeflash/lsp/server.py b/codeflash/lsp/server.py index a018786a8..e89ac7ea4 100644 --- a/codeflash/lsp/server.py +++ b/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. From a0ff82f530e3adaff4cb426e91f73f825c0f788e Mon Sep 17 00:00:00 2001 From: mohammed Date: Fri, 11 Jul 2025 21:07:14 +0300 Subject: [PATCH 2/3] fix: clear cache before checking the api key to pick the latest one, and enable codeflash debug messages in lsp --- codeflash/code_utils/shell_utils.py | 2 +- codeflash/lsp/beta.py | 41 ++++++++++++++++++----------- codeflash/lsp/server.py | 3 ++- 3 files changed, 28 insertions(+), 18 deletions(-) diff --git a/codeflash/code_utils/shell_utils.py b/codeflash/code_utils/shell_utils.py index a0198c366..30a5aadaa 100644 --- a/codeflash/code_utils/shell_utils.py +++ b/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]: diff --git a/codeflash/lsp/beta.py b/codeflash/lsp/beta.py index 10335c550..9d554e1d8 100644 --- a/codeflash/lsp/beta.py +++ b/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"} diff --git a/codeflash/lsp/server.py b/codeflash/lsp/server.py index e89ac7ea4..030574a8d 100644 --- a/codeflash/lsp/server.py +++ b/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) From abeb01e4d51e4e6788685cd8fd26e7b2876fdc84 Mon Sep 17 00:00:00 2001 From: mohammed Date: Wed, 6 Aug 2025 16:08:45 +0300 Subject: [PATCH 3/3] better error handling and read api key from rc only --- codeflash/api/cfapi.py | 2 +- codeflash/code_utils/env_utils.py | 2 +- codeflash/lsp/beta.py | 4 ++++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/codeflash/api/cfapi.py b/codeflash/api/cfapi.py index 958375d88..7948c4970 100644 --- a/codeflash/api/cfapi.py +++ b/codeflash/api/cfapi.py @@ -101,7 +101,7 @@ def get_user_id() -> Optional[str]: console.print(f"[bold red]{msg}[/bold red]") if console.quiet: # lsp logger.debug(msg) - return None + return f"Error: {msg}" sys.exit(1) return userid diff --git a/codeflash/code_utils/env_utils.py b/codeflash/code_utils/env_utils.py index 476edebdf..5692e1a82 100644 --- a/codeflash/code_utils/env_utils.py +++ b/codeflash/code_utils/env_utils.py @@ -36,7 +36,7 @@ def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool = def get_codeflash_api_key() -> str: 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") + api_key = read_api_key_from_shell_config() else: api_key = os.environ.get("CODEFLASH_API_KEY") or read_api_key_from_shell_config() diff --git a/codeflash/lsp/beta.py b/codeflash/lsp/beta.py index 9d554e1d8..2e5492a0c 100644 --- a/codeflash/lsp/beta.py +++ b/codeflash/lsp/beta.py @@ -130,6 +130,10 @@ def _initialize_optimizer_if_valid(server: CodeflashLanguageServer) -> dict[str, if user_id is None: return {"status": "error", "message": "api key not found or invalid"} + if user_id.startswith("Error: "): + error_msg = user_id[7:] + return {"status": "error", "message": error_msg} + from codeflash.optimization.optimizer import Optimizer server.optimizer = Optimizer(server.args)