[LSP] Message Id tag

[mohammedahmed18]

User description

send a message id with the lsp logs for the client to know when it will receive a specific type of message

currently only telling the client when it start testing a candidate, and when it finds the best candidate, later we can add more to it.

this can be used in the future to implement steps or timeline of the optimization which we discussed before

---

PR Type

Enhancement

---

Description

• Add LSP message IDs to logs

• Propagate IDs through code and text messages

• Make task_id optional in execution context

• Parse and serialize message_id in LSP messages

---

Diagram Walkthrough

flowchart LR
  console["console.code_print adds lsp_message_id"] -- "passes to" --> lsp_msg["LspCodeMessage/LspTextMessage.message_id"]
  lsp_logger["lsp_logger.extract_tags parses id:"] -- "sets" --> tags["LspMessageTags.message_id"]
  tags -- "serialized via" --> serialize["LspTextMessage.serialize() with message_id"]
  optimizer["function_optimizer emits IDs for candidate/best"] -- "uses" --> enum["LSPMessageId Enum"]
  beta["lsp.beta execution_context"] -- "uses getattr optional task_id" --> execctx["execution_context"]

Loading

File Walkthrough

Relevant files

Enhancement

console.py code_print supports LSP message IDs                                            codeflash/cli_cmds/console.py Add lsp_message_id param to code_print. Forward message_id to LspCodeMessage. +9/-2      beta.py Optional task_id in execution context usage                            codeflash/lsp/beta.py Use getattr(params, "task_id", None) in two handlers. Makes execution_context robust to missing task_id. +2/-2      lsp_logger.py Logger parses and forwards LSP message IDs                              codeflash/lsp/lsp_logger.py Add message_id to LspMessageTags. Parse id: tags into LSPMessageId. Include message_id when serializing LspTextMessage. +27/-21  lsp_message.py LSP message model gains message_id enum                                    codeflash/lsp/lsp_message.py Introduce LSPMessageId enum. Add message_id to LspMessage base. +8/-0      function_optimizer.py Emit message IDs for candidates and best                                  codeflash/optimization/function_optimizer.py Import LSPMessageId. Attach CANDIDATE ID to candidate code prints. Attach BEST_CANDIDATE ID to best candidate print. +7/-2

---

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 2 🔵🔵⚪⚪⚪

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Possible Issue Assigning message_tags.message_id = LSPMessageId(...).value assumes all ids after the 'id:' prefix map to enum values; unexpected ids will raise ValueError during enum construction. Consider defensive handling (try/except or validating membership) to avoid crashing logging. def extract_tags(msg: str) -> tuple[LspMessageTags, str]: delimiter = "|" first_delim_idx = msg.find(delimiter) if first_delim_idx != -1 and msg.count(delimiter) == 1: tags_str = msg[:first_delim_idx] content = msg[first_delim_idx + 1 :] tags = {tag.strip() for tag in tags_str.split(",")} message_tags = LspMessageTags() # manually check and set to avoid repeated membership tests for tag in tags: if tag.startswith(message_id_prefix): message_tags.message_id = LSPMessageId(tag[len(message_id_prefix) :]).value elif tag == "lsp": message_tags.lsp = True elif tag == "!lsp": message_tags.not_lsp = True elif tag == "force_lsp": message_tags.force_lsp = True elif tag == "loading": message_tags.loading = True elif tag == "highlight": message_tags.highlight = True elif tag == "h1": message_tags.h1 = True elif tag == "h2": message_tags.h2 = True elif tag == "h3": message_tags.h3 = True elif tag == "h4": message_tags.h4 = True return message_tags, content Type Consistency LspMessage.message_id is typed as Optional[str] while callers pass LSPMessageId.value in some places and LSPMessageId is also imported elsewhere. Consider standardizing on Optional[LSPMessageId] or consistently using str to avoid confusion and accidental enum leaks. # allow the client to know which message it is receiving class LSPMessageId(enum.Enum): BEST_CANDIDATE = "best_candidate" CANDIDATE = "candidate" @dataclass class LspMessage: # to show a loading indicator if the operation is taking time like generating candidates or tests takes_time: bool = False message_id: Optional[str] = None def _loop_through(self, obj: Any) -> Any: # noqa: ANN401 API Change Ripple New parameter lsp_message_id was added to code_print; verify all internal callers are updated or have defaults. Also ensure non-LSP path ignores the param safely and that static typing (if any) is updated. def code_print( code_str: str, file_name: Optional[str] = None, function_name: Optional[str] = None, lsp_message_id: Optional[str] = None, ) -> None: if is_LSP_enabled(): lsp_log( LspCodeMessage(code=code_str, file_name=file_name, function_name=function_name, message_id=lsp_message_id) ) return """Print code with syntax highlighting.""" from rich.syntax import Syntax

[github-actions]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Handle unknown message IDs safely Guard against invalid or unknown message_id values to avoid ValueError when constructing LSPMessageId. Fall back to treating the tag content as a raw string if it doesn't match the enum. This keeps logging resilient to unexpected tags. codeflash/lsp/lsp_logger.py [59-79] for tag in tags: if tag.startswith(message_id_prefix): - message_tags.message_id = LSPMessageId(tag[len(message_id_prefix) :]).value + raw_id = tag[len(message_id_prefix) :] + try: + message_tags.message_id = LSPMessageId(raw_id).value + except ValueError: + # accept arbitrary IDs without failing enum conversion + message_tags.message_id = raw_id elif tag == "lsp": message_tags.lsp = True elif tag == "!lsp": message_tags.not_lsp = True elif tag == "force_lsp": message_tags.force_lsp = True elif tag == "loading": message_tags.loading = True elif tag == "highlight": message_tags.highlight = True elif tag == "h1": message_tags.h1 = True elif tag == "h2": message_tags.h2 = True elif tag == "h3": message_tags.h3 = True elif tag == "h4": message_tags.h4 = True Suggestion importance[1-10]: 7 __ Why: This guards enum construction against invalid message_id values, preventing a potential ValueError at runtime. It's directly applicable to the new loop and improves robustness without altering behavior for valid IDs.	Medium
General	Sanitize optional message ID Validate or normalize lsp_message_id before forwarding to the logger to avoid emitting empty or whitespace-only IDs. Only attach message_id when a meaningful value is provided. This prevents ambiguous client-side routing. codeflash/cli_cmds/console.py [90-92] +normalized_id = lsp_message_id.strip() if isinstance(lsp_message_id, str) else None lsp_log( - LspCodeMessage(code=code_str, file_name=file_name, function_name=function_name, message_id=lsp_message_id) + LspCodeMessage(code=code_str, file_name=file_name, function_name=function_name, message_id=normalized_id or None) ) Suggestion importance[1-10]: 6 __ Why: Trimming and omitting empty message_id avoids emitting meaningless identifiers and is a safe, low-risk improvement aligned with the new parameter usage.	Low
	Ensure non-empty task context If task_id is required by execution_context, passing None may break tracing or cause collisions. Provide a stable fallback (e.g., derive from params) to ensure context is always populated. This prevents missing-context issues in logs and cancellations. codeflash/lsp/beta.py [356] -with execution_context(task_id=getattr(params, "task_id", None)): +task_id = getattr(params, "task_id", None) or getattr(params, "functionName", None) or getattr(params, "textDocument", None) and getattr(getattr(params, "textDocument", None), "uri", None) +with execution_context(task_id=task_id): Suggestion importance[1-10]: 5 __ Why: Providing a fallback for task_id may help maintain tracing when it's missing, but correctness depends on execution_context requirements and the suitability of fallbacks; impact is moderate and somewhat speculative.	Low