[LSP] stderr verbose logs for the thought process (#718)

* lsp silent logs

* override other log methods and log serialized lsp messages

* send the module root to the lsp client

* lsp messages

* code print over lsp

* more enhancements

* more enhancements

* cf optimization

* log tags for lsp

* better markdown support for lsp message logging

* simple markdown table

* force lsp log (tag)

* fixes for cli console logs

* small fixes

* small fix

* it should work this time

* prevent worktree log in lsp

* logging enhancement

* file name for best candidate

* reminder

* lsp logs formatting and small fixes

* typo

* fixes for the api key and the lsp gracefull shutdown

---------

Co-authored-by: Sarthak Agarwal <sarthak.saga@gmail.com>

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -133,7 +133,7 @@ def optimize_python_code(  # noqa: D417
             "repo_name": git_repo_name,
         }
 
-        logger.info("Generating optimized candidates…")
+        logger.info("!lsp|Generating optimized candidates…")
         console.rule()
         try:
             response = self.make_ai_service_request("/optimize", payload=payload, timeout=600)
@@ -144,10 +144,10 @@ def optimize_python_code(  # noqa: D417
 
         if response.status_code == 200:
             optimizations_json = response.json()["optimizations"]
-            logger.info(f"Generated {len(optimizations_json)} candidate optimizations.")
+            logger.info(f"!lsp|Generated {len(optimizations_json)} candidate optimizations.")
             console.rule()
             end_time = time.perf_counter()
-            logger.debug(f"Generating optimizations took {end_time - start_time:.2f} seconds.")
+            logger.debug(f"!lsp|Generating possible optimizations took {end_time - start_time:.2f} seconds.")
             return self._get_valid_candidates(optimizations_json)
         try:
             error = response.json()["error"]
@@ -194,7 +194,6 @@ def optimize_python_code_line_profiler(  # noqa: D417
             "lsp_mode": is_LSP_enabled(),
         }
 
-        logger.info("Generating optimized candidates…")
         console.rule()
         if line_profiler_results == "":
             logger.info("No LineProfiler results were provided, Skipping optimization.")
@@ -209,7 +208,9 @@ def optimize_python_code_line_profiler(  # noqa: D417
 
         if response.status_code == 200:
             optimizations_json = response.json()["optimizations"]
-            logger.info(f"Generated {len(optimizations_json)} candidate optimizations using line profiler information.")
+            logger.info(
+                f"!lsp|Generated {len(optimizations_json)} candidate optimizations using line profiler information."
+            )
             console.rule()
             return self._get_valid_candidates(optimizations_json)
         try:
@@ -331,7 +332,7 @@ def get_new_explanation(  # noqa: D417
             "original_explanation": original_explanation,
             "dependency_code": dependency_code,
         }
-        logger.info("Generating explanation")
+        logger.info("loading|Generating explanation")
         console.rule()
         try:
             response = self.make_ai_service_request("/explain", payload=payload, timeout=60)
@@ -376,7 +377,7 @@ def generate_ranking(  # noqa: D417
             "optimization_ids": optimization_ids,
             "python_version": platform.python_version(),
         }
-        logger.info("Generating ranking")
+        logger.info("loading|Generating ranking")
         console.rule()
         try:
             response = self.make_ai_service_request("/rank", payload=payload, timeout=60)

codeflash/api/cfapi.py

@@ -40,6 +40,7 @@ def make_cfapi_request(
     payload: dict[str, Any] | None = None,
     extra_headers: dict[str, str] | None = None,
     *,
+    api_key: str | None = None,
     suppress_errors: bool = False,
 ) -> Response:
     """Make an HTTP request using the specified method, URL, headers, and JSON payload.
@@ -51,7 +52,7 @@ def make_cfapi_request(
     :return: The response object from the API.
     """
     url = f"{CFAPI_BASE_URL}/cfapi{endpoint}"
-    cfapi_headers = {"Authorization": f"Bearer {get_codeflash_api_key()}"}
+    cfapi_headers = {"Authorization": f"Bearer {api_key or get_codeflash_api_key()}"}
     if extra_headers:
         cfapi_headers.update(extra_headers)
     try:
@@ -83,15 +84,17 @@ def make_cfapi_request(
 
 
 @lru_cache(maxsize=1)
-def get_user_id() -> Optional[str]:
+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.
 
     :return: The userid or None if the request fails.
     """
     if not ensure_codeflash_api_key():
         return None
 
-    response = make_cfapi_request(endpoint="/cli-get-user", method="GET", extra_headers={"cli_version": __version__})
+    response = make_cfapi_request(
+        endpoint="/cli-get-user", method="GET", extra_headers={"cli_version": __version__}, api_key=api_key
+    )
     if response.status_code == 200:
         if "min_version" not in response.text:
             return response.text

codeflash/cli_cmds/console.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
 import logging
-import os
 from contextlib import contextmanager
 from itertools import cycle
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from rich.console import Console
 from rich.logging import RichHandler
@@ -20,17 +19,22 @@
 
 from codeflash.cli_cmds.console_constants import SPINNER_TYPES
 from codeflash.cli_cmds.logging_config import BARE_LOGGING_FORMAT
+from codeflash.lsp.helpers import is_LSP_enabled
+from codeflash.lsp.lsp_logger import enhanced_log
+from codeflash.lsp.lsp_message import LspCodeMessage, LspTextMessage
 
 if TYPE_CHECKING:
     from collections.abc import Generator
 
     from rich.progress import TaskID
 
+    from codeflash.lsp.lsp_message import LspMessage
+
 DEBUG_MODE = logging.getLogger().getEffectiveLevel() == logging.DEBUG
 
 console = Console()
 
-if os.getenv("CODEFLASH_LSP"):
+if is_LSP_enabled():
     console.quiet = True
 
 logging.basicConfig(
@@ -42,6 +46,24 @@
 logger = logging.getLogger("rich")
 logging.getLogger("parso").setLevel(logging.WARNING)
 
+# override the logger to reformat the messages for the lsp
+for level in ("info", "debug", "warning", "error"):
+    real_fn = getattr(logger, level)
+    setattr(
+        logger,
+        level,
+        lambda msg, *args, _real_fn=real_fn, _level=level, **kwargs: enhanced_log(
+            msg, _real_fn, _level, *args, **kwargs
+        ),
+    )
+
+
+def lsp_log(message: LspMessage) -> None:
+    if not is_LSP_enabled():
+        return
+    json_msg = message.serialize()
+    logger.info(json_msg)
+
 
 def paneled_text(
     text: str, panel_args: dict[str, str | bool] | None = None, text_args: dict[str, str] | None = None
@@ -58,7 +80,10 @@ def paneled_text(
     console.print(panel)
 
 
-def code_print(code_str: str) -> None:
+def code_print(code_str: str, file_name: Optional[str] = None, function_name: Optional[str] = None) -> None:
+    if is_LSP_enabled():
+        lsp_log(LspCodeMessage(code=code_str, file_name=file_name, function_name=function_name))
+        return
     """Print code with syntax highlighting."""
     from rich.syntax import Syntax
 
@@ -79,6 +104,11 @@ def progress_bar(
     If revert_to_print is True, falls back to printing a single logger.info message
     instead of showing a progress bar.
     """
+    if is_LSP_enabled():
+        lsp_log(LspTextMessage(text=message, takes_time=True))
+        yield
+        return
+
     if revert_to_print:
         logger.info(message)
 

codeflash/code_utils/git_worktree_utils.py

@@ -62,7 +62,7 @@ def create_detached_worktree(module_root: Path) -> Optional[Path]:
     )
 
     if not uni_diff_text.strip():
-        logger.info("No uncommitted changes to copy to worktree.")
+        logger.info("!lsp|No uncommitted changes to copy to worktree.")
         return worktree_dir
 
     # Write the diff to a temporary file

codeflash/discovery/functions_to_optimize.py

@@ -173,15 +173,15 @@ def get_functions_to_optimize(
     with warnings.catch_warnings():
         warnings.simplefilter(action="ignore", category=SyntaxWarning)
         if optimize_all:
-            logger.info("Finding all functions in the module '%s'…", optimize_all)
+            logger.info("!lsp|Finding all functions in the module '%s'…", optimize_all)
             console.rule()
             functions = get_all_files_and_functions(Path(optimize_all))
         elif replay_test:
             functions, trace_file_path = get_all_replay_test_functions(
                 replay_test=replay_test, test_cfg=test_cfg, project_root_path=project_root
             )
         elif file is not None:
-            logger.info("Finding all functions in the file '%s'…", file)
+            logger.info("!lsp|Finding all functions in the file '%s'…", file)
             console.rule()
             functions = find_all_functions_in_file(file)
             if only_get_this_function is not None:
@@ -219,7 +219,7 @@ def get_functions_to_optimize(
             functions, test_cfg.tests_root, ignore_paths, project_root, module_root, previous_checkpoint_functions
         )
 
-        logger.info(f"Found {functions_count} function{'s' if functions_count > 1 else ''} to optimize")
+        logger.info(f"!lsp|Found {functions_count} function{'s' if functions_count > 1 else ''} to optimize")
         if optimize_all:
             three_min_in_ns = int(1.8e11)
             console.rule()

codeflash/lsp/beta.py

@@ -4,13 +4,14 @@
 import os
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 import git
 from pygls import uris
 
 from codeflash.api.cfapi import get_codeflash_api_key, get_user_id
 from codeflash.cli_cmds.cli import process_pyproject_config
+from codeflash.cli_cmds.console import code_print
 from codeflash.code_utils.git_worktree_utils import (
     create_diff_patch_from_worktree,
     get_patches_metadata,
@@ -103,6 +104,8 @@ def get_optimizable_functions(
 ) -> dict[str, list[str]]:
     file_path = Path(uris.to_fs_path(params.textDocument.uri))
     server.show_message_log(f"Getting optimizable functions for: {file_path}", "Info")
+    if not server.optimizer:
+        return {"status": "error", "message": "optimizer not initialized"}
 
     server.optimizer.args.file = file_path
     server.optimizer.args.function = None  # Always get ALL functions, not just one
@@ -157,20 +160,6 @@ def initialize_function_optimization(
     return {"functionName": params.functionName, "status": "success"}
 
 
-@server.feature("discoverFunctionTests")
-def discover_function_tests(server: CodeflashLanguageServer, params: FunctionOptimizationParams) -> dict[str, str]:
-    fto = server.optimizer.current_function_being_optimized
-    optimizable_funcs = {fto.file_path: [fto]}
-
-    devnull_writer = open(os.devnull, "w")  # noqa
-    with contextlib.redirect_stdout(devnull_writer):
-        function_to_tests, num_discovered_tests = server.optimizer.discover_tests(optimizable_funcs)
-
-    server.optimizer.discovered_tests = function_to_tests
-
-    return {"functionName": params.functionName, "status": "success", "discovered_tests": num_discovered_tests}
-
-
 @server.feature("validateProject")
 def validate_project(server: CodeflashLanguageServer, _params: FunctionOptimizationParams) -> dict[str, str]:
     from codeflash.cli_cmds.cmd_init import is_valid_pyproject_toml
@@ -194,11 +183,13 @@ def validate_project(server: CodeflashLanguageServer, _params: FunctionOptimizat
     except Exception:
         return {"status": "error", "message": "Repository has no commits (unborn HEAD)"}
 
-    return {"status": "success"}
+    return {"status": "success", "moduleRoot": args.module_root}
 
 
-def _initialize_optimizer_if_api_key_is_valid(server: CodeflashLanguageServer) -> dict[str, str]:
-    user_id = get_user_id()
+def _initialize_optimizer_if_api_key_is_valid(
+    server: CodeflashLanguageServer, api_key: Optional[str] = None
+) -> dict[str, str]:
+    user_id = get_user_id(api_key=api_key)
     if user_id is None:
         return {"status": "error", "message": "api key not found or invalid"}
 
@@ -237,19 +228,19 @@ def provide_api_key(server: CodeflashLanguageServer, params: ProvideApiKeyParams
         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_api_key_is_valid(server)
+        init_result = _initialize_optimizer_if_api_key_is_valid(server, api_key)
         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"]}
+        user_id = init_result["user_id"]
+        result = save_api_key_to_rc(api_key)
+        if not is_successful(result):
+            return {"status": "error", "message": result.failure()}
+        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"}
 
@@ -300,6 +291,12 @@ def perform_function_optimization(  # noqa: PLR0911
             }
 
         module_prep_result = server.optimizer.prepare_module_for_optimization(current_function.file_path)
+        if not module_prep_result:
+            return {
+                "functionName": params.functionName,
+                "status": "error",
+                "message": "Failed to prepare module for optimization",
+            }
 
         validated_original_code, original_module_ast = module_prep_result
 
@@ -308,7 +305,7 @@ def perform_function_optimization(  # noqa: PLR0911
             function_to_optimize_source_code=validated_original_code[current_function.file_path].source_code,
             original_module_ast=original_module_ast,
             original_module_path=current_function.file_path,
-            function_to_tests=server.optimizer.discovered_tests or {},
+            function_to_tests={},
         )
 
         server.optimizer.current_function_optimizer = function_optimizer
@@ -321,6 +318,19 @@ def perform_function_optimization(  # noqa: PLR0911
 
         should_run_experiment, code_context, original_helper_code = initialization_result.unwrap()
 
+        code_print(
+            code_context.read_writable_code.flat,
+            file_name=current_function.file_path,
+            function_name=current_function.function_name,
+        )
+
+        optimizable_funcs = {current_function.file_path: [current_function]}
+
+        devnull_writer = open(os.devnull, "w")  # noqa
+        with contextlib.redirect_stdout(devnull_writer):
+            function_to_tests, num_discovered_tests = server.optimizer.discover_tests(optimizable_funcs)
+            function_optimizer.function_to_tests = function_to_tests
+
         test_setup_result = function_optimizer.generate_and_instrument_tests(
             code_context, should_run_experiment=should_run_experiment
         )