Merge branch 'saga4/claude_mcp_and_subagent' of https://github.com/codeflash-ai/codeflash into saga4/claude_mcp_and_subagent

Author: saga4

.pre-commit-config.yaml

@@ -1,7 +1,8 @@
 repos:
-  - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.11.0"
-    hooks:
-      - id: ruff
-        args: [--fix, --exit-non-zero-on-fix, --config=pyproject.toml]
-      - id: ruff-format
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.12.7
+  hooks:
+    # Run the linter.
+    - id: ruff-check
+    # Run the formatter.
+    - id: ruff-format

README.md

@@ -65,6 +65,13 @@ For detailed installation and usage instructions, visit our documentation at [do
 
 https://github.com/user-attachments/assets/38f44f4e-be1c-4f84-8db9-63d5ee3e61e5
 
+- Optiming a workflow end to end automatically with `codeflash optimize`
+
+
+https://github.com/user-attachments/assets/355ba295-eb5a-453a-8968-7fb35c70d16c
+
+
+
 ## Support
 
 Join our community for support and discussions. If you have any questions, feel free to reach out to us using one of the following methods:

SECURITY.md

@@ -0,0 +1,19 @@
+# Security Policy
+
+This document outlines Codeflash's vulnerability disclosure policy. For more information about Codeflash's approach to security, please visit [codeflash.ai/security](https://www.codeflash.ai/security).
+
+## Supported Versions
+
+Since Codeflash is moving quickly, we can only commit to fixing security issues for the latest version of codeflash client.
+If a vulnerability is discovered in our backend, we will release the fix for all the users.
+
+## Reporting a Vulnerability
+
+
+Please do not report security vulnerabilities through public GitHub issues.
+
+Instead, please report them to our [GitHub Security page](https://github.com/codeflash-ai/codeflash/security). If you prefer to submit one without using GitHub, you can also email us at [email protected].
+
+We commit to acknowledging vulnerability reports immediately, and will work to fix active vulnerabilities as soon as we can. We will publish resolved vulnerabilities in the form of security advisories on our GitHub security page. Critical incidents will be communicated both on the GitHub security page and via email to all affected users.
+
+We appreciate your help in making Codeflash more secure for everyone. Thank you for your support and responsible disclosure.

codeflash/api/cfapi.py

@@ -101,6 +101,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 f"Error: {msg}"
                 sys.exit(1)
             return userid
 

codeflash/benchmarking/instrument_codeflash_trace.py

@@ -32,12 +32,14 @@ def leave_ClassDef(
     def visit_ClassDef(self, node: ClassDef) -> Optional[bool]:
         if self.class_name:  # Don't go into nested class
             return False
-        self.class_name = node.name.value  # noqa: RET503
+        self.class_name = node.name.value
+        return None
 
     def visit_FunctionDef(self, node: FunctionDef) -> Optional[bool]:
         if self.function_name:  # Don't go into nested function
             return False
-        self.function_name = node.name.value  # noqa: RET503
+        self.function_name = node.name.value
+        return None
 
     def leave_FunctionDef(self, original_node: FunctionDef, updated_node: FunctionDef) -> FunctionDef:
         if self.function_name == original_node.name.value:

codeflash/cli_cmds/cmd_init.py

@@ -38,7 +38,7 @@
     from argparse import Namespace
 
 CODEFLASH_LOGO: str = (
-    f"{LF}"  # noqa: ISC003
+    f"{LF}"
     r"                   _          ___  _               _     " + f"{LF}"
     r"                  | |        / __)| |             | |    " + f"{LF}"
     r"  ____   ___    _ | |  ____ | |__ | |  ____   ___ | | _  " + f"{LF}"

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()
+    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/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,6 +8,8 @@
 
 from pygls import uris
 
+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
 
@@ -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,53 @@ def discover_function_tests(server: CodeflashLanguageServer, params: FunctionOpt
     return {"functionName": params.functionName, "status": "success", "discovered_tests": num_discovered_tests}
 
 
+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"}
+
+    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)
+    return {"status": "success", "user_id": user_id}
+
+
+@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"}
+
+
+@server.feature("provideApiKey")
+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()}
+
+        # 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"}
+
+        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"}
+
+
 @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
+        # 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.
@@ -70,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)