Integrate LSP for debugging (#450)

* headless IDE subclass

* finish headless_ide methods

* feat: :sparkles: headless mode running with config flag

* work on debugging

* python lsp support

* more lsp+debugging work

* refactor: :safety_vest: safely load LSP

* test: :white_check_mark: testing steps in headless mode

* refactor: :clown_face: cleanup subprocesses

* fix: :bug: handle data: [DONE] from Together

4 files changed, 374 insertions, 18 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index bae82739..8ac7241d 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -17,7 +17,10 @@ from ..libs.util.paths import getSavedContextGroupsPath
 from ..libs.util.queue import AsyncSubscriptionQueue
 from ..libs.util.strings import remove_quotes_and_escapes
 from ..libs.util.telemetry import posthog_logger
-from ..libs.util.traceback_parsers import get_javascript_traceback, get_python_traceback
+from ..libs.util.traceback.traceback_parsers import (
+    get_javascript_traceback,
+    get_python_traceback,
+)
 from ..models.filesystem import RangeInFileWithContents
 from ..models.filesystem_edit import FileEditWithFullContents
 from ..models.main import ContinueBaseModel
@@ -32,6 +35,8 @@ from ..plugins.steps.core.core import (
 )
 from ..plugins.steps.on_traceback import DefaultOnTracebackStep
 from ..server.ide_protocol import AbstractIdeProtocolServer
+from ..server.meilisearch_server import stop_meilisearch
+from .config import ContinueConfig
 from .context import ContextManager
 from .main import (
     Context,
@@ -97,8 +102,12 @@ class Autopilot(ContinueBaseModel):
 
     started: bool = False
 
-    async def start(self, full_state: Optional[FullState] = None):
-        self.continue_sdk = await ContinueSDK.create(self)
+    async def start(
+        self,
+        full_state: Optional[FullState] = None,
+        config: Optional[ContinueConfig] = None,
+    ):
+        self.continue_sdk = await ContinueSDK.create(self, config=config)
         if override_policy := self.continue_sdk.config.policy_override:
             self.policy = override_policy
 
@@ -134,6 +143,11 @@ class Autopilot(ContinueBaseModel):
 
         self.started = True
 
+    async def cleanup(self):
+        if self.continue_sdk.lsp is not None:
+            await self.continue_sdk.lsp.stop()
+        stop_meilisearch()
+
     class Config:
         arbitrary_types_allowed = True
         keep_untouched = (cached_property,)
diff --git a/continuedev/src/continuedev/core/config.py b/continuedev/src/continuedev/core/config.py
index 62e9c690..68b2b17d 100644
--- a/continuedev/src/continuedev/core/config.py
+++ b/continuedev/src/continuedev/core/config.py
@@ -59,3 +59,14 @@ class ContinueConfig(BaseModel):
     @validator("temperature", pre=True)
     def temperature_validator(cls, v):
         return max(0.0, min(1.0, v))
+
+    @staticmethod
+    def from_filepath(filepath: str) -> "ContinueConfig":
+        # Use importlib to load the config file config.py at the given path
+        import importlib.util
+
+        spec = importlib.util.spec_from_file_location("config", filepath)
+        config = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(config)
+
+        return config.config
diff --git a/continuedev/src/continuedev/core/lsp.py b/continuedev/src/continuedev/core/lsp.py
new file mode 100644
index 00000000..5c1f9989
--- /dev/null
+++ b/continuedev/src/continuedev/core/lsp.py
@@ -0,0 +1,310 @@
+import os
+import socket
+import subprocess
+import threading
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+from ..libs.lspclient.json_rpc_endpoint import JsonRpcEndpoint
+from ..libs.lspclient.lsp_client import LspClient
+from ..libs.lspclient.lsp_endpoint import LspEndpoint
+from ..libs.lspclient.lsp_structs import Position as LspPosition
+from ..libs.lspclient.lsp_structs import SymbolInformation, TextDocumentIdentifier
+from ..models.filesystem import RangeInFile
+from ..models.main import Position, Range
+
+
+class ReadPipe(threading.Thread):
+    def __init__(self, pipe):
+        threading.Thread.__init__(self)
+        self.pipe = pipe
+
+    def run(self):
+        line = self.pipe.readline().decode("utf-8")
+        while line:
+            print(line)
+            line = self.pipe.readline().decode("utf-8")
+
+
+class SocketFileWrapper:
+    def __init__(self, sockfile):
+        self.sockfile = sockfile
+
+    def write(self, data):
+        if isinstance(data, bytes):
+            data = data.decode("utf-8").replace("\r\n", "\n")
+        return self.sockfile.write(data)
+
+    def read(self, size=-1):
+        data = self.sockfile.read(size)
+        if isinstance(data, str):
+            data = data.replace("\n", "\r\n").encode("utf-8")
+        return data
+
+    def readline(self, size=-1):
+        data = self.sockfile.readline(size)
+        if isinstance(data, str):
+            data = data.replace("\n", "\r\n").encode("utf-8")
+        return data
+
+    def flush(self):
+        return self.sockfile.flush()
+
+    def close(self):
+        return self.sockfile.close()
+
+
+def create_json_rpc_endpoint(use_subprocess: Optional[str] = None):
+    if use_subprocess is None:
+        # Connect to the server
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.connect(("localhost", 8080))
+
+        # Create a file-like object from the socket
+        sockfile = s.makefile("rw")
+        wrapped_sockfile = SocketFileWrapper(sockfile)
+        return JsonRpcEndpoint(wrapped_sockfile, wrapped_sockfile), None
+
+    else:
+        pyls_cmd = use_subprocess.split()
+        p = subprocess.Popen(
+            pyls_cmd,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+        read_pipe = ReadPipe(p.stderr)
+        read_pipe.start()
+        return JsonRpcEndpoint(p.stdin, p.stdout), p
+
+
+def filename_to_uri(filename: str) -> str:
+    return f"file://{filename}"
+
+
+def uri_to_filename(uri: str) -> str:
+    if uri.startswith("file://"):
+        return uri.lstrip("file://")
+    else:
+        return uri
+
+
+def create_lsp_client(workspace_dir: str, use_subprocess: Optional[str] = None):
+    json_rpc_endpoint, process = create_json_rpc_endpoint(use_subprocess=use_subprocess)
+    lsp_endpoint = LspEndpoint(json_rpc_endpoint)
+    lsp_client = LspClient(lsp_endpoint)
+    capabilities = {
+        "textDocument": {
+            "codeAction": {"dynamicRegistration": True},
+            "codeLens": {"dynamicRegistration": True},
+            "colorProvider": {"dynamicRegistration": True},
+            "completion": {
+                "completionItem": {
+                    "commitCharactersSupport": True,
+                    "documentationFormat": ["markdown", "plaintext"],
+                    "snippetSupport": True,
+                },
+                "completionItemKind": {
+                    "valueSet": [
+                        1,
+                        2,
+                        3,
+                        4,
+                        5,
+                        6,
+                        7,
+                        8,
+                        9,
+                        10,
+                        11,
+                        12,
+                        13,
+                        14,
+                        15,
+                        16,
+                        17,
+                        18,
+                        19,
+                        20,
+                        21,
+                        22,
+                        23,
+                        24,
+                        25,
+                    ]
+                },
+                "contextSupport": True,
+                "dynamicRegistration": True,
+            },
+            "definition": {"dynamicRegistration": True},
+            "documentHighlight": {"dynamicRegistration": True},
+            "documentLink": {"dynamicRegistration": True},
+            "documentSymbol": {
+                "dynamicRegistration": True,
+                "symbolKind": {
+                    "valueSet": [
+                        1,
+                        2,
+                        3,
+                        4,
+                        5,
+                        6,
+                        7,
+                        8,
+                        9,
+                        10,
+                        11,
+                        12,
+                        13,
+                        14,
+                        15,
+                        16,
+                        17,
+                        18,
+                        19,
+                        20,
+                        21,
+                        22,
+                        23,
+                        24,
+                        25,
+                        26,
+                    ]
+                },
+            },
+            "formatting": {"dynamicRegistration": True},
+            "hover": {
+                "contentFormat": ["markdown", "plaintext"],
+                "dynamicRegistration": True,
+            },
+            "implementation": {"dynamicRegistration": True},
+            "onTypeFormatting": {"dynamicRegistration": True},
+            "publishDiagnostics": {"relatedInformation": True},
+            "rangeFormatting": {"dynamicRegistration": True},
+            "references": {"dynamicRegistration": True},
+            "rename": {"dynamicRegistration": True},
+            "signatureHelp": {
+                "dynamicRegistration": True,
+                "signatureInformation": {
+                    "documentationFormat": ["markdown", "plaintext"]
+                },
+            },
+            "synchronization": {
+                "didSave": True,
+                "dynamicRegistration": True,
+                "willSave": True,
+                "willSaveWaitUntil": True,
+            },
+            "typeDefinition": {"dynamicRegistration": True},
+        },
+        "workspace": {
+            "applyEdit": True,
+            "configuration": True,
+            "didChangeConfiguration": {"dynamicRegistration": True},
+            "didChangeWatchedFiles": {"dynamicRegistration": True},
+            "executeCommand": {"dynamicRegistration": True},
+            "symbol": {
+                "dynamicRegistration": True,
+                "symbolKind": {
+                    "valueSet": [
+                        1,
+                        2,
+                        3,
+                        4,
+                        5,
+                        6,
+                        7,
+                        8,
+                        9,
+                        10,
+                        11,
+                        12,
+                        13,
+                        14,
+                        15,
+                        16,
+                        17,
+                        18,
+                        19,
+                        20,
+                        21,
+                        22,
+                        23,
+                        24,
+                        25,
+                        26,
+                    ]
+                },
+            },
+            "workspaceEdit": {"documentChanges": True},
+            "workspaceFolders": True,
+        },
+    }
+    root_uri = filename_to_uri(workspace_dir)
+    dir_name = os.path.basename(workspace_dir)
+    workspace_folders = [{"name": dir_name, "uri": root_uri}]
+    lsp_client.initialize(
+        None,
+        None,
+        root_uri,
+        None,
+        capabilities,
+        "off",
+        workspace_folders,
+    )
+    lsp_client.initialized()
+    return lsp_client, process
+
+
+class ContinueLSPClient(BaseModel):
+    workspace_dir: str
+    lsp_client: LspClient = None
+    use_subprocess: Optional[str] = None
+    lsp_process: Optional[subprocess.Popen] = None
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def dict(self, **kwargs):
+        original_dict = super().dict(**kwargs)
+        original_dict.pop("lsp_client", None)
+        return original_dict
+
+    async def start(self):
+        self.lsp_client, self.lsp_process = create_lsp_client(
+            self.workspace_dir, use_subprocess=self.use_subprocess
+        )
+
+    async def stop(self):
+        self.lsp_client.shutdown()
+        self.lsp_client.exit()
+        if self.lsp_process is not None:
+            self.lsp_process.terminate()
+            self.lsp_process.wait()
+            self.lsp_process = None
+
+    def goto_definition(self, position: Position, filename: str) -> List[RangeInFile]:
+        response = self.lsp_client.definition(
+            TextDocumentIdentifier(filename_to_uri(filename)),
+            LspPosition(position.line, position.character),
+        )
+        return [
+            RangeInFile(
+                filepath=uri_to_filename(x.uri),
+                range=Range.from_shorthand(
+                    x.range.start.line,
+                    x.range.start.character,
+                    x.range.end.line,
+                    x.range.end.character,
+                ),
+            )
+            for x in response
+        ]
+
+    def get_symbols(self, filepath: str) -> List[SymbolInformation]:
+        response = self.lsp_client.documentSymbol(
+            TextDocumentIdentifier(filename_to_uri(filepath))
+        )
+
+        return response
diff --git a/continuedev/src/continuedev/core/sdk.py b/continuedev/src/continuedev/core/sdk.py
index 9ff6612c..7dca600d 100644
--- a/continuedev/src/continuedev/core/sdk.py
+++ b/continuedev/src/continuedev/core/sdk.py
@@ -1,8 +1,9 @@
 import os
 import traceback
-from typing import Coroutine, Union
+from typing import Coroutine, List, Optional, Union
 
 from ..libs.llm import LLM
+from ..libs.util.create_async_task import create_async_task
 from ..libs.util.logging import logger
 from ..libs.util.paths import getConfigFilePath
 from ..libs.util.telemetry import posthog_logger
@@ -16,11 +17,18 @@ from ..models.filesystem_edit import (
     FileSystemEdit,
 )
 from ..models.main import Range
-from ..plugins.steps.core.core import *
-from ..plugins.steps.core.core import DefaultModelEditCodeStep
+from ..plugins.steps.core.core import (
+    DefaultModelEditCodeStep,
+    FileSystemEditStep,
+    MessageStep,
+    RangeInFileWithContents,
+    ShellCommandsStep,
+    WaitForUserConfirmationStep,
+)
 from ..server.ide_protocol import AbstractIdeProtocolServer
 from .abstract_sdk import AbstractContinueSDK
 from .config import ContinueConfig
+from .lsp import ContinueLSPClient
 from .main import (
     ChatMessage,
     Context,
@@ -42,6 +50,7 @@ class ContinueSDK(AbstractContinueSDK):
 
     ide: AbstractIdeProtocolServer
     models: Models
+    lsp: Optional[ContinueLSPClient] = None
     context: Context
     config: ContinueConfig
     __autopilot: Autopilot
@@ -52,13 +61,14 @@ class ContinueSDK(AbstractContinueSDK):
         self.context = autopilot.context
 
     @classmethod
-    async def create(cls, autopilot: Autopilot) -> "ContinueSDK":
+    async def create(
+        cls, autopilot: Autopilot, config: Optional[ContinueConfig] = None
+    ) -> "ContinueSDK":
         sdk = ContinueSDK(autopilot)
         autopilot.continue_sdk = sdk
 
         try:
-            config = sdk._load_config_dot_py()
-            sdk.config = config
+            sdk.config = config or sdk._load_config_dot_py()
         except Exception as e:
             logger.error(f"Failed to load config.py: {traceback.format_exception(e)}")
 
@@ -78,9 +88,26 @@ class ContinueSDK(AbstractContinueSDK):
             )
             await sdk.ide.setFileOpen(getConfigFilePath())
 
+        # Start models
         sdk.models = sdk.config.models
         await sdk.models.start(sdk)
 
+        # Start LSP
+        async def start_lsp():
+            try:
+                sdk.lsp = ContinueLSPClient(
+                    workspace_dir=sdk.ide.workspace_directory,
+                    use_subprocess="python3.10 -m pylsp",
+                )
+                await sdk.lsp.start()
+            except:
+                logger.warning("Failed to start LSP client", exc_info=True)
+                sdk.lsp = None
+
+        create_async_task(
+            start_lsp(), on_error=lambda e: logger.error("Failed to setup LSP: %s", e)
+        )
+
         # When the config is loaded, setup posthog logger
         posthog_logger.setup(sdk.ide.unique_id, sdk.config.allow_anonymous_telemetry)
 
@@ -207,19 +234,13 @@ class ContinueSDK(AbstractContinueSDK):
     _last_valid_config: ContinueConfig = None
 
     def _load_config_dot_py(self) -> ContinueConfig:
-        # Use importlib to load the config file config.py at the given path
         path = getConfigFilePath()
-
-        import importlib.util
-
-        spec = importlib.util.spec_from_file_location("config", path)
-        config = importlib.util.module_from_spec(spec)
-        spec.loader.exec_module(config)
-        self._last_valid_config = config.config
+        config = ContinueConfig.from_filepath(path)
+        self._last_valid_config = config
 
         logger.debug("Loaded Continue config file from %s", path)
 
-        return config.config
+        return config
 
     def get_code_context(
         self, only_editing: bool = False