[Fix] Circular Import Bug

codeflash/lsp/beta.py

@@ -29,6 +29,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
 
@@ -336,12 +337,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 typing import TYPE_CHECKING
 
 from lsprotocol.types import LogMessageParams, MessageType
@@ -25,10 +24,6 @@ def __init__(self, name: str, version: str, protocol_cls: type[LanguageServerPro
         self.optimizer: Optimizer | None = None
         self.args = None
         self.current_optimization_init_result: tuple[bool, CodeOptimizationContext, dict[Path, str]] | None = None
-        self.execution_context_vars: contextvars.ContextVar[dict[str, str]] = contextvars.ContextVar(
-            "execution_context_vars",
-            default={},  # noqa: B039
-        )
 
     def prepare_optimizer_arguments(self, config_file: Path) -> None:
         from codeflash.cli_cmds.cli import parse_args