better markdown support for lsp message logging

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -133,7 +133,6 @@ def optimize_python_code(  # noqa: D417
             "repo_name": git_repo_name,
         }
 
-        logger.info("Generating optimized candidates…")
         console.rule()
         try:
             response = self.make_ai_service_request("/optimize", payload=payload, timeout=600)
@@ -144,8 +143,6 @@ 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.")
-            console.rule()
             end_time = time.perf_counter()
             logger.debug(f"Generating optimizations took {end_time - start_time:.2f} seconds.")
             return self._get_valid_candidates(optimizations_json)
@@ -194,7 +191,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
             "lsp_mode": is_LSP_enabled(),
         }
 
-        logger.info("Generating optimized candidates…")
+        logger.info("loading|tags|Generating optimized candidates using line profiler information…")
         console.rule()
         if line_profiler_results == "":
             logger.info("No LineProfiler results were provided, Skipping optimization.")
@@ -209,7 +206,6 @@ 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.")
             console.rule()
             return self._get_valid_candidates(optimizations_json)
         try:

codeflash/cli_cmds/console.py

@@ -49,7 +49,13 @@
 # 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, **kwargs: enhanced_log(msg, _real_fn, *args, **kwargs))
+    setattr(
+        logger,
+        level,
+        lambda msg, _real_fn=real_fn, _level=level, *args, **kwargs: enhanced_log(
+            msg, _real_fn, _level, *args, **kwargs
+        ),
+    )
 
 
 def lsp_log(message: LspMessage) -> None:

codeflash/lsp/beta.py

@@ -158,20 +158,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
@@ -301,6 +287,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
 
@@ -309,7 +301,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
@@ -328,6 +320,13 @@ def perform_function_optimization(  # noqa: PLR0911
             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
         )

codeflash/lsp/helpers.py

@@ -2,17 +2,44 @@
 import re
 from functools import lru_cache
 
-_double_quote_pat = re.compile(r'"(.*?)"')
+from rich.tree import Tree
+
+from codeflash.models.test_type import TestType
 
+_double_quote_pat = re.compile(r'"(.*?)"')
 _single_quote_pat = re.compile(r"'(.*?)'")
+worktree_path_regex = re.compile(r'\/[^"]*worktrees\/[^"]\S*')
 
 
 @lru_cache(maxsize=1)
 def is_LSP_enabled() -> bool:
     return os.getenv("CODEFLASH_LSP", default="false").lower() == "true"
 
 
-worktree_path_regex = re.compile(r'\/[^"]*worktrees\/[^"]\S*')
+def tree_to_markdown(tree: Tree, level: int = 0) -> str:
+    """Convert a rich Tree into a Markdown bullet list."""
+    indent = "  " * level
+    if level == 0:
+        lines: list[str] = [f"{indent}### {tree.label}"]
+    else:
+        lines: list[str] = [f"{indent}- {tree.label}"]
+    for child in tree.children:
+        lines.extend(tree_to_markdown(child, level + 1).splitlines())
+    return "\n".join(lines)
+
+
+def report_to_markdown_table(report: dict[TestType, dict[str, int]], title: str) -> str:
+    lines = ["| Test Type | Passed ✅ | Failed ❌ |", "|-----------|--------|--------|"]
+    for test_type in TestType:
+        if test_type is TestType.INIT_STATE_TEST:
+            continue
+        passed = report[test_type]["passed"]
+        failed = report[test_type]["failed"]
+        # if passed == 0 and failed == 0:
+        #     continue
+        lines.append(f"| {test_type.to_name()} | {passed} | {failed} |")
+    table = "\n".join(lines)
+    return f"### {title}\n{table}"
 
 
 def simplify_worktree_paths(msg: str, highlight: bool = True) -> str:  # noqa: FBT001, FBT002

codeflash/lsp/lsp_logger.py

@@ -13,6 +13,23 @@ class LspMessageTags:
     not_lsp: bool = False
     loading: bool = False
 
+    h1: bool = False
+    h2: bool = False
+    h3: bool = False
+    h4: bool = False
+
+
+def add_heading_tags(msg: str, tags: LspMessageTags) -> str:
+    if tags.h1:
+        return "# " + msg
+    if tags.h2:
+        return "## " + msg
+    if tags.h3:
+        return "### " + msg
+    if tags.h4:
+        return "#### " + msg
+    return msg
+
 
 def extract_tags(msg: str) -> tuple[LspMessageTags, str]:
     if not isinstance(msg, str):
@@ -26,12 +43,29 @@ def extract_tags(msg: str) -> tuple[LspMessageTags, str]:
             message_tags.not_lsp = True
         if "loading" in tags:
             message_tags.loading = True
+        if "h1" in tags:
+            message_tags.h1 = True
+        if "h2" in tags:
+            message_tags.h2 = True
+        if "h3" in tags:
+            message_tags.h3 = True
+        if "h4" in tags:
+            message_tags.h4 = True
         return message_tags, parts[1]
 
     return LspMessageTags(), msg
 
 
-def enhanced_log(msg: str, actual_log_fn: Callable[[str, Any, Any], None], *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
+supported_lsp_log_levels = ("info", "debug")
+
+
+def enhanced_log(
+    msg: str,
+    actual_log_fn: Callable[[str, Any, Any], None],
+    level: str,
+    *args: Any,  # noqa: ANN401
+    **kwargs: Any,  # noqa: ANN401
+) -> None:
     lsp_enabled = is_LSP_enabled()
     if not lsp_enabled or not isinstance(msg, str):
         actual_log_fn(msg, *args, **kwargs)
@@ -52,10 +86,11 @@ def enhanced_log(msg: str, actual_log_fn: Callable[[str, Any, Any], None], *args
         return
 
     #### LSP mode ####
-    if tags.not_lsp:
+    if tags.not_lsp or level not in supported_lsp_log_levels:
         return
 
     if is_normal_text_message:
+        clean_msg = add_heading_tags(clean_msg, tags)
         clean_msg = LspTextMessage(text=clean_msg, takes_time=tags.loading).serialize()
 
     actual_log_fn(clean_msg, *args, **kwargs)

codeflash/lsp/lsp_message.py

@@ -31,19 +31,12 @@ def type(self) -> str:
         raise NotImplementedError
 
     def serialize(self) -> str:
-        if isinstance(self, LspTextMessage):
-            self.text = simplify_worktree_paths(self.text)
-            self.text = replace_quotes_with_backticks(self.text)
-        if isinstance(self, LspCodeMessage):
-            self.file_name = simplify_worktree_paths(str(self.file_name), highlight=False)
-        if isinstance(self, LspMarkdownMessage):
-            self.markdown = simplify_worktree_paths(self.markdown)
-            self.markdown = replace_quotes_with_backticks(self.markdown)
-
         data = self._loop_through(asdict(self))
         # Important: keep type as the first key, for making it easy and fast for the client to know if this is a lsp message before parsing it
         ordered = {"type": self.type(), **data}
-        return json.dumps(ordered)
+        return (
+            json.dumps(ordered) + "\u241f"
+        )  # \u241F is the message delimiter becuase it can be more than one message sent over the same message
 
 
 @dataclass
@@ -53,6 +46,11 @@ class LspTextMessage(LspMessage):
     def type(self) -> str:
         return "text"
 
+    def serialize(self) -> str:
+        self.text = simplify_worktree_paths(self.text)
+        self.text = replace_quotes_with_backticks(self.text)
+        return super().serialize()
+
 
 @dataclass
 class LspCodeMessage(LspMessage):
@@ -70,6 +68,7 @@ def serialize(self) -> str:
         self.lines_count = code_lines_length
         if code_lines_length > max_code_lines_before_collapse:
             self.collapsed = True
+        self.file_name = simplify_worktree_paths(str(self.file_name), highlight=False)
         return super().serialize()
 
 
@@ -80,10 +79,7 @@ class LspMarkdownMessage(LspMessage):
     def type(self) -> str:
         return "markdown"
 
-
-@dataclass
-class LspStatsMessage(LspMessage):
-    stats: dict[str, Any]
-
-    def type(self) -> str:
-        return "stats"
+    def serialize(self) -> str:
+        self.markdown = simplify_worktree_paths(self.markdown)
+        self.markdown = replace_quotes_with_backticks(self.markdown)
+        return super().serialize()

codeflash/models/models.py

@@ -5,7 +5,10 @@
 
 from rich.tree import Tree
 
-from codeflash.cli_cmds.console import DEBUG_MODE
+from codeflash.cli_cmds.console import DEBUG_MODE, lsp_log
+from codeflash.lsp.helpers import is_LSP_enabled, report_to_markdown_table
+from codeflash.lsp.lsp_message import LspMarkdownMessage
+from codeflash.models.test_type import TestType
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
@@ -482,27 +485,6 @@ class VerificationType(str, Enum):
     INIT_STATE_HELPER = "init_state_helper"  # Correctness verification for helper class instance attributes after init
 
 
-class TestType(Enum):
-    EXISTING_UNIT_TEST = 1
-    INSPIRED_REGRESSION = 2
-    GENERATED_REGRESSION = 3
-    REPLAY_TEST = 4
-    CONCOLIC_COVERAGE_TEST = 5
-    INIT_STATE_TEST = 6
-
-    def to_name(self) -> str:
-        if self is TestType.INIT_STATE_TEST:
-            return ""
-        names = {
-            TestType.EXISTING_UNIT_TEST: "⚙️ Existing Unit Tests",
-            TestType.INSPIRED_REGRESSION: "🎨 Inspired Regression Tests",
-            TestType.GENERATED_REGRESSION: "🌀 Generated Regression Tests",
-            TestType.REPLAY_TEST: "⏪ Replay Tests",
-            TestType.CONCOLIC_COVERAGE_TEST: "🔎 Concolic Coverage Tests",
-        }
-        return names[self]
-
-
 @dataclass(frozen=True)
 class InvocationId:
     test_module_path: str  # The fully qualified name of the test module
@@ -644,6 +626,13 @@ def report_to_string(report: dict[TestType, dict[str, int]]) -> str:
     @staticmethod
     def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree:
         tree = Tree(title)
+
+        if is_LSP_enabled():
+            # Build markdown table
+            markdown = report_to_markdown_table(report, title)
+            lsp_log(LspMarkdownMessage(markdown=markdown))
+            return tree
+
         for test_type in TestType:
             if test_type is TestType.INIT_STATE_TEST:
                 continue

codeflash/models/test_type.py

@@ -0,0 +1,22 @@
+from enum import Enum
+
+
+class TestType(Enum):
+    EXISTING_UNIT_TEST = 1
+    INSPIRED_REGRESSION = 2
+    GENERATED_REGRESSION = 3
+    REPLAY_TEST = 4
+    CONCOLIC_COVERAGE_TEST = 5
+    INIT_STATE_TEST = 6
+
+    def to_name(self) -> str:
+        if self is TestType.INIT_STATE_TEST:
+            return ""
+        names = {
+            TestType.EXISTING_UNIT_TEST: "⚙️ Existing Unit Tests",
+            TestType.INSPIRED_REGRESSION: "🎨 Inspired Regression Tests",
+            TestType.GENERATED_REGRESSION: "🌀 Generated Regression Tests",
+            TestType.REPLAY_TEST: "⏪ Replay Tests",
+            TestType.CONCOLIC_COVERAGE_TEST: "🔎 Concolic Coverage Tests",
+        }
+        return names[self]