lsp messages

Author: mohammedahmed18

codeflash/lsp/beta.py

@@ -194,7 +194,7 @@ def validate_project(server: CodeflashLanguageServer, _params: FunctionOptimizat
     except Exception:
         return {"status": "error", "message": "Repository has no commits (unborn HEAD)"}
 
-    return {"status": "success", "module_root": args.module_root}
+    return {"status": "success", "moduleRoot": args.module_root}
 
 
 def _initialize_optimizer_if_api_key_is_valid(server: CodeflashLanguageServer) -> dict[str, str]:

codeflash/lsp/helpers.py

@@ -2,6 +2,8 @@
 from functools import lru_cache
 from typing import Any, Callable
 
+skip_lsp_log_prefix = "!lsp:"
+
 
 @lru_cache(maxsize=1)
 def is_LSP_enabled() -> bool:
@@ -10,18 +12,19 @@ def is_LSP_enabled() -> bool:
 
 def enhanced_log(msg: str, actual_log_fn: Callable[[str, Any, Any], None], *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
     lsp_enabled = is_LSP_enabled()
+    str_msg = isinstance(msg, str)
+    skip_lsp_log = str_msg and msg.strip().startswith(skip_lsp_log_prefix)
+
+    if skip_lsp_log:
+        msg = msg[len(skip_lsp_log_prefix) :]
 
-    # normal cli moded
+    # normal cli mode
     if not lsp_enabled:
         actual_log_fn(msg, *args, **kwargs)
         return
 
     #### LSP mode ####
-    if type(msg) != str:  # noqa: E721
-        return
-
-    if msg.startswith("Nonzero return code"):
-        # skip logging the failed tests msg to the client
+    if skip_lsp_log or not str_msg:
         return
 
     actual_log_fn(f"::::{msg}", *args, **kwargs)

codeflash/lsp/lsp_message.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any, Optional
+
+json_primitive_types = (str, float, int, bool)
+
+
+@dataclass
+class LspMessage:
+    takes_time: bool = field(
+        default=False, kw_only=True
+    )  # to show a loading indicator if the operation is taking time like generating candidates or tests
+
+    def _loop_through(self, obj: Any) -> Any:  # noqa: ANN401
+        if isinstance(obj, list):
+            return [self._loop_through(i) for i in obj]
+        if isinstance(obj, dict):
+            return {k: self._loop_through(v) for k, v in obj.items()}
+        if isinstance(obj, json_primitive_types) or obj is None:
+            return obj
+        if isinstance(obj, Path):
+            return obj.as_posix()
+        return str(obj)
+
+    def type(self) -> str:
+        raise NotImplementedError
+
+    def serialize(self) -> str:
+        data = asdict(self)
+        data["type"] = self.type()
+        return json.dumps(data)
+
+
+@dataclass
+class LspTextMessage(LspMessage):
+    text: str
+
+    def type(self) -> str:
+        return "text"
+
+
+@dataclass
+class LspCodeMessage(LspMessage):
+    code: str
+    path: Optional[Path] = None
+    function_name: Optional[str] = None
+
+    def type(self) -> str:
+        return "code"
+
+
+@dataclass
+class LspMarkdownMessage(LspMessage):
+    markdown: str
+
+    def type(self) -> str:
+        return "markdown"
+
+
+@dataclass
+class LspStatsMessage(LspMessage):
+    stats: dict[str, Any]
+
+    def type(self) -> str:
+        return "stats"

codeflash/optimization/function_optimizer.py

@@ -1673,7 +1673,7 @@ def run_and_parse_tests(
             return TestResults(), None
         if run_result.returncode != 0 and testing_type == TestingMode.BEHAVIOR:
             logger.debug(
-                f"Nonzero return code {run_result.returncode} when running tests in "
+                f"!lsp:Nonzero return code {run_result.returncode} when running tests in "
                 f"{', '.join([str(f.instrumented_behavior_file_path) for f in test_files.test_files])}.\n"
                 f"stdout: {run_result.stdout}\n"
                 f"stderr: {run_result.stderr}\n"