feat: :sparkles: LSP connection over websockets

2 files changed, 273 insertions, 268 deletions

diff --git a/continuedev/src/continuedev/core/lsp.py b/continuedev/src/continuedev/core/lsp.py
index 86923fb7..181eea2e 100644
--- a/continuedev/src/continuedev/core/lsp.py
+++ b/continuedev/src/continuedev/core/lsp.py
@@ -1,282 +1,272 @@
 import asyncio
-import os
-import socket
-import subprocess
 import threading
 from typing import List, Optional
 
+import aiohttp
 from pydantic import BaseModel
-from pylsp.python_lsp import PythonLSPServer, start_tcp_lang_server
+from pylsp.python_lsp import PythonLSPServer, start_ws_lang_server
 
-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 ..libs.util.logging import logger
 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")
+def filepath_to_uri(filename: str) -> str:
+    return f"file://{filename}"
 
 
-class SocketFileWrapper:
-    def __init__(self, sockfile):
-        self.sockfile = sockfile
+def uri_to_filepath(uri: str) -> str:
+    if uri.startswith("file://"):
+        return uri.lstrip("file://")
+    else:
+        return uri
 
-    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
+PORT = 8099
 
-    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()
+class LSPClient:
+    def __init__(self, host: str, port: int, workspace_paths: List[str]):
+        self.host = host
+        self.port = port
+        self.session = aiohttp.ClientSession()
+        self.next_id = 0
+        self.workspace_paths = workspace_paths
 
-    def close(self):
-        return self.sockfile.close()
+    async def connect(self):
+        print("Connecting")
+        self.ws = await self.session.ws_connect(f"ws://{self.host}:{self.port}/")
+        print("Connected")
 
+    async def send(self, data):
+        await self.ws.send_json(data)
 
-async def create_json_rpc_endpoint(use_subprocess: Optional[str] = None):
-    if use_subprocess is None:
-        try:
-            threading.Thread(
-                target=start_tcp_lang_server,
-                args=("localhost", 8080, False, PythonLSPServer),
-            ).start()
-        except Exception as e:
-            logger.warning("Could not start TCP server: %s", e)
+    async def recv(self):
+        return await self.ws.receive_json()
 
-        await asyncio.sleep(2)
+    async def close(self):
+        await self.ws.close()
+        await self.session.close()
 
-        # Connect to the server
-        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        s.connect(("localhost", 8080))
+    async def call_method(self, method_name, **kwargs):
+        body = {
+            "jsonrpc": "2.0",
+            "id": self.next_id,
+            "method": method_name,
+            "params": kwargs,
+        }
+        self.next_id += 1
+        await self.send(body)
+        response = await self.recv()
+        return response
 
-        # Create a file-like object from the socket
-        sockfile = s.makefile("rw")
-        wrapped_sockfile = SocketFileWrapper(sockfile)
-        return JsonRpcEndpoint(wrapped_sockfile, wrapped_sockfile), None
+    async def initialize(self):
+        initialization_args = {
+            "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,
+                },
+            },
+            "processId": 1234,
+            "rootPath": None,
+            "rootUri": filepath_to_uri(self.workspace_paths[0]),
+            "initializationOptions": {},
+            "trace": "off",
+            "workspaceFolders": [
+                {
+                    "uri": filepath_to_uri(workspacePath),
+                    "name": workspacePath.split("/")[-1],
+                }
+                for workspacePath in self.workspace_paths
+            ],
+        }
+        return await self.call_method("initialize", **initialization_args)
+
+    async def goto_definition(self, filepath: str, position: Position):
+        return await self.call_method(
+            "textDocument/definition",
+            textDocument={"uri": filepath_to_uri(filepath)},
+            position=position.dict(),
+        )
 
-    else:
-        pyls_cmd = use_subprocess.split()
-        p = subprocess.Popen(
-            pyls_cmd,
-            stdin=subprocess.PIPE,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
+    async def document_symbol(self, filepath: str):
+        return await self.call_method(
+            "textDocument/documentSymbol",
+            textDocument={"uri": filepath_to_uri(filepath)},
         )
-        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}"
+async def start_language_server() -> threading.Thread:
+    try:
+        thread = threading.Thread(
+            target=start_ws_lang_server,
+            args=(PORT, False, PythonLSPServer),
+        )
+        thread.daemon = True
+        thread.start()
 
+    except Exception as e:
+        logger.warning("Could not start TCP server: %s", e)
 
-def uri_to_filename(uri: str) -> str:
-    if uri.startswith("file://"):
-        return uri.lstrip("file://")
-    else:
-        return uri
+    await asyncio.sleep(2)
 
+    return thread
 
-async def create_lsp_client(workspace_dir: str, use_subprocess: Optional[str] = None):
-    json_rpc_endpoint, process = await 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 DocumentSymbol(BaseModel):
+    name: str
+    containerName: Optional[str] = None
+    kind: int
+    location: RangeInFile
 
 
 class ContinueLSPClient(BaseModel):
     workspace_dir: str
-    lsp_client: LspClient = None
-    use_subprocess: Optional[str] = None
-    lsp_process: Optional[subprocess.Popen] = None
+
+    lsp_client: LSPClient = None
+    lsp_thread: Optional[threading.Thread] = None
 
     class Config:
         arbitrary_types_allowed = True
@@ -287,26 +277,26 @@ class ContinueLSPClient(BaseModel):
         return original_dict
 
     async def start(self):
-        self.lsp_client, self.lsp_process = await create_lsp_client(
-            self.workspace_dir, use_subprocess=self.use_subprocess
-        )
+        self.lsp_thread = await start_language_server()
+        self.lsp_client = LSPClient("localhost", PORT, [self.workspace_dir])
+        await self.lsp_client.connect()
+        await self.lsp_client.initialize()
 
     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),
+        await self.lsp_client.close()
+        if self.lsp_thread:
+            self.lsp_thread.join()
+
+    async def goto_definition(
+        self, position: Position, filename: str
+    ) -> List[RangeInFile]:
+        response = self.lsp_client.goto_definition(
+            filename,
+            position,
         )
         return [
             RangeInFile(
-                filepath=uri_to_filename(x.uri),
+                filepath=uri_to_filepath(x.uri),
                 range=Range.from_shorthand(
                     x.range.start.line,
                     x.range.start.character,
@@ -317,9 +307,22 @@ class ContinueLSPClient(BaseModel):
             for x in response
         ]
 
-    def get_symbols(self, filepath: str) -> List[SymbolInformation]:
-        response = self.lsp_client.documentSymbol(
-            TextDocumentIdentifier(filename_to_uri(filepath))
-        )
-
-        return response
+    async def document_symbol(self, filepath: str) -> List:
+        response = await self.lsp_client.document_symbol(filepath)
+        return [
+            DocumentSymbol(
+                name=x["name"],
+                containerName=x["containerName"],
+                kind=x["kind"],
+                location=RangeInFile(
+                    filepath=uri_to_filepath(x["location"]["uri"]),
+                    range=Range.from_shorthand(
+                        x["location"]["range"]["start"]["line"],
+                        x["location"]["range"]["start"]["character"],
+                        x["location"]["range"]["end"]["line"],
+                        x["location"]["range"]["end"]["character"],
+                    ),
+                ),
+            )
+            for x in response["result"]
+        ]
diff --git a/continuedev/src/continuedev/core/sdk.py b/continuedev/src/continuedev/core/sdk.py
index 658848c8..de209114 100644
--- a/continuedev/src/continuedev/core/sdk.py
+++ b/continuedev/src/continuedev/core/sdk.py
@@ -2,6 +2,8 @@ import os
 import traceback
 from typing import Coroutine, List, Optional, Union
 
+from ..libs.util.create_async_task import create_async_task
+
 from ..libs.llm import LLM
 from ..libs.util.devdata import dev_data_logger
 from ..libs.util.logging import logger
@@ -103,9 +105,9 @@ class ContinueSDK(AbstractContinueSDK):
                 logger.warning(f"Failed to start LSP client: {e}", exc_info=True)
                 sdk.lsp = None
 
-        # create_async_task(
-        #     start_lsp(), on_error=lambda e: logger.error("Failed to setup LSP: %s", e)
-        # )
+        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)