Merge branch 'main' of github.com:codeflash-ai/codeflash into lsp/demo-optimization

Author: mohammedahmed18

codeflash/api/cfapi.py

@@ -1,3 +1,5 @@
+"""Module for interacting with the Codeflash API."""
+
 from __future__ import annotations
 
 import json
@@ -14,6 +16,7 @@
 from pydantic.json import pydantic_encoder
 
 from codeflash.cli_cmds.console import console, logger
+from codeflash.code_utils.code_utils import exit_with_message
 from codeflash.code_utils.env_utils import ensure_codeflash_api_key, get_codeflash_api_key, get_pr_number
 from codeflash.code_utils.git_utils import get_current_branch, get_repo_owner_and_name
 from codeflash.github.PrComment import FileDiffContent, PrComment
@@ -101,13 +104,18 @@ def make_cfapi_request(
 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.
 
+    :param api_key: The API key to use. If None, uses get_codeflash_api_key().
     :return: The userid or None if the request fails.
     """
     if not api_key and not ensure_codeflash_api_key():
         return None
 
     response = make_cfapi_request(
-        endpoint="/cli-get-user", method="GET", extra_headers={"cli_version": __version__}, api_key=api_key
+        endpoint="/cli-get-user",
+        method="GET",
+        extra_headers={"cli_version": __version__},
+        api_key=api_key,
+        suppress_errors=True,
     )
     if response.status_code == 200:
         if "min_version" not in response.text:
@@ -128,6 +136,20 @@ def get_user_id(api_key: Optional[str] = None) -> Optional[str]:
         logger.error("Failed to retrieve userid from the response.")
         return None
 
+    # Handle 403 (Invalid API key) - exit with error message
+    if response.status_code == 403:
+        msg = (
+            "Invalid Codeflash API key. The API key you provided is not valid.\n"
+            "Please generate a new one at https://app.codeflash.ai/app/apikeys ,\n"
+            "then set it as a CODEFLASH_API_KEY environment variable.\n"
+            "For more information, refer to the documentation at \n"
+            "https://docs.codeflash.ai/optimizing-with-codeflash/codeflash-github-actions#manual-setup\n"
+            "or\n"
+            "https://docs.codeflash.ai/optimizing-with-codeflash/codeflash-github-actions#automated-setup-recommended"
+        )
+        exit_with_message(msg, error_on_exit=True)
+
+    # For other errors, log and return None (backward compatibility)
     logger.error(f"Failed to look up your userid; is your CF API key valid? ({response.reason})")
     return None
 

codeflash/code_utils/github_utils.py

@@ -4,7 +4,7 @@
 
 from codeflash.api.cfapi import is_github_app_installed_on_repo
 from codeflash.cli_cmds.cli_common import apologize_and_exit
-from codeflash.cli_cmds.console import logger
+from codeflash.cli_cmds.console import paneled_text
 from codeflash.code_utils.compat import LF
 from codeflash.code_utils.git_utils import get_repo_owner_and_name
 
@@ -18,15 +18,20 @@ def get_github_secrets_page_url(repo: Optional[Repo] = None) -> str:
 
 
 def require_github_app_or_exit(owner: str, repo: str) -> None:
-    if not is_github_app_installed_on_repo(owner, repo):
-        logger.error(
-            f"It looks like the Codeflash GitHub App is not installed on the repository {owner}/{repo} or the GitHub"
-            f" account linked to your CODEFLASH_API_KEY does not have access to the repository {owner}/{repo}.{LF}"
-            "Before continuing, please install the Codeflash GitHub App on your repository by visiting "
-            f"https://github.com/apps/codeflash-ai/installations/select_target{LF}"
+    # Suppress low-level HTTP error logging to avoid duplicate logs; we present a friendly panel instead
+    if not is_github_app_installed_on_repo(owner, repo, suppress_errors=True):
+        # Show a clear, user-friendly panel instead of raw error logs
+        message = (
+            f"It looks like the Codeflash GitHub App is not installed on the repository {owner}/{repo} "
+            f"or the GitHub account linked to your CODEFLASH_API_KEY does not have access to the repository {owner}/{repo}.{LF}{LF}"
+            "To continue, install the Codeflash GitHub App on your repository:"
+            f"{LF}https://github.com/apps/codeflash-ai/installations/select_target{LF}{LF}"
+            "Tip: If you want to find optimizations without opening PRs, run Codeflash with the --no-pr flag."
         )
-        logger.error(
-            f"Note: if you want to find optimizations without opening PRs, you can run Codeflash with the --no-pr flag.{LF}"
+        paneled_text(
+            message,
+            panel_args={"title": "GitHub App Required", "border_style": "red", "expand": False},
+            text_args={"style": "bold red"},
         )
         apologize_and_exit()
 

codeflash/lsp/beta.py

@@ -31,6 +31,7 @@
     get_functions_within_git_diff,
 )
 from codeflash.either import is_successful
+from codeflash.lsp.context import execution_context_vars
 from codeflash.lsp.features.perform_optimization import get_cancelled_reponse, sync_perform_optimization
 from codeflash.lsp.server import CodeflashLanguageServer, CodeflashLanguageServerProtocol
 
@@ -341,12 +342,12 @@ def provide_api_key(params: ProvideApiKeyParams) -> dict[str, str]:
 def execution_context(**kwargs: str) -> None:
     """Temporarily set context values for the current async task."""
     # Create a fresh copy per use
-    current = {**server.execution_context_vars.get(), **kwargs}
-    token = server.execution_context_vars.set(current)
+    current = {**execution_context_vars.get(), **kwargs}
+    token = execution_context_vars.set(current)
     try:
         yield
     finally:
-        server.execution_context_vars.reset(token)
+        execution_context_vars.reset(token)
 
 
 @server.feature("cleanupCurrentOptimizerSession")

codeflash/lsp/context.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+import contextvars
+
+# Shared execution context for tracking task IDs and other metadata
+execution_context_vars: contextvars.ContextVar[dict[str, str]] = contextvars.ContextVar(
+    "execution_context_vars",
+    default={},  # noqa: B039
+)

codeflash/lsp/lsp_message.py

@@ -44,9 +44,9 @@ def type(self) -> str:
     def serialize(self) -> str:
         if not is_LSP_enabled():
             return ""
-        from codeflash.lsp.beta import server
+        from codeflash.lsp.context import execution_context_vars
 
-        execution_ctx = server.execution_context_vars.get()
+        execution_ctx = execution_context_vars.get()
         current_task_id = execution_ctx.get("task_id", None)
         data = self._loop_through(asdict(self))
         ordered = {"type": self.type(), "task_id": current_task_id, **data}

codeflash/lsp/server.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import contextvars
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -29,11 +28,7 @@ def __init__(self, name: str, version: str, protocol_cls: type[LanguageServerPro
         self.initialized: bool = False
         self.optimizer: Optimizer | None = None
         self.args = None
-        self.current_optimization_init_result: InitializationResultT | None = None
-        self.execution_context_vars: contextvars.ContextVar[dict[str, str]] = contextvars.ContextVar(
-            "execution_context_vars",
-            default={},  # noqa: B039
-        )
+        self.current_optimization_init_result: tuple[bool, CodeOptimizationContext, dict[Path, str]] | None = None
 
     def prepare_optimizer_arguments(self, config_file: Path) -> None:
         from codeflash.cli_cmds.cli import parse_args