Refactor helper (#481)

* feat: :sparkles: add stop_tokens option to LLM

* work on refactoring in headless mode

* feat: :sparkles: headless mode refactors

* chore: :fire: remove test.py

1 files changed, 113 insertions, 24 deletions

diff --git a/continuedev/src/continuedev/core/lsp.py b/continuedev/src/continuedev/core/lsp.py
index 181eea2e..0c906b22 100644
--- a/continuedev/src/continuedev/core/lsp.py
+++ b/continuedev/src/continuedev/core/lsp.py
@@ -1,6 +1,6 @@
 import asyncio
 import threading
-from typing import List, Optional
+from typing import List, Literal, Optional
 
 import aiohttp
 from pydantic import BaseModel
@@ -9,6 +9,7 @@ from pylsp.python_lsp import PythonLSPServer, start_ws_lang_server
 from ..libs.util.logging import logger
 from ..models.filesystem import RangeInFile
 from ..models.main import Position, Range
+from ..server.meilisearch_server import kill_proc
 
 
 def filepath_to_uri(filename: str) -> str:
@@ -17,7 +18,7 @@ def filepath_to_uri(filename: str) -> str:
 
 def uri_to_filepath(uri: str) -> str:
     if uri.startswith("file://"):
-        return uri.lstrip("file://")
+        return uri[7:]
     else:
         return uri
 
@@ -26,6 +27,9 @@ PORT = 8099
 
 
 class LSPClient:
+    ready: bool = False
+    lock: asyncio.Lock = asyncio.Lock()
+
     def __init__(self, host: str, port: int, workspace_paths: List[str]):
         self.host = host
         self.port = port
@@ -37,12 +41,18 @@ class LSPClient:
         print("Connecting")
         self.ws = await self.session.ws_connect(f"ws://{self.host}:{self.port}/")
         print("Connected")
+        self.ready = True
 
     async def send(self, data):
         await self.ws.send_json(data)
 
     async def recv(self):
-        return await self.ws.receive_json()
+        await self.lock.acquire()
+
+        try:
+            return await self.ws.receive_json()
+        finally:
+            self.lock.release()
 
     async def close(self):
         await self.ws.close()
@@ -237,9 +247,27 @@ class LSPClient:
             textDocument={"uri": filepath_to_uri(filepath)},
         )
 
+    async def find_references(
+        self, filepath: str, position: Position, include_declaration: bool = False
+    ):
+        return await self.call_method(
+            "textDocument/references",
+            textDocument={"uri": filepath_to_uri(filepath)},
+            position=position.dict(),
+            context={"includeDeclaration": include_declaration},
+        )
+
+    async def folding_range(self, filepath: str):
+        response = await self.call_method(
+            "textDocument/foldingRange",
+            textDocument={"uri": filepath_to_uri(filepath)},
+        )
+        return response["result"]
+
 
 async def start_language_server() -> threading.Thread:
     try:
+        kill_proc(PORT)
         thread = threading.Thread(
             target=start_ws_lang_server,
             args=(PORT, False, PythonLSPServer),
@@ -262,12 +290,23 @@ class DocumentSymbol(BaseModel):
     location: RangeInFile
 
 
+class FoldingRange(BaseModel):
+    range: Range
+    kind: Optional[Literal["comment", "imports", "region"]] = None
+
+
 class ContinueLSPClient(BaseModel):
     workspace_dir: str
 
     lsp_client: LSPClient = None
     lsp_thread: Optional[threading.Thread] = None
 
+    @property
+    def ready(self):
+        if self.lsp_client is None:
+            return False
+        return self.lsp_client.ready
+
     class Config:
         arbitrary_types_allowed = True
 
@@ -287,6 +326,17 @@ class ContinueLSPClient(BaseModel):
         if self.lsp_thread:
             self.lsp_thread.join()
 
+    def location_to_range_in_file(self, location):
+        return RangeInFile(
+            filepath=uri_to_filepath(location["uri"]),
+            range=Range.from_shorthand(
+                location["range"]["start"]["line"],
+                location["range"]["start"]["character"],
+                location["range"]["end"]["line"],
+                location["range"]["end"]["character"],
+            ),
+        )
+
     async def goto_definition(
         self, position: Position, filename: str
     ) -> List[RangeInFile]:
@@ -294,18 +344,17 @@ class ContinueLSPClient(BaseModel):
             filename,
             position,
         )
-        return [
-            RangeInFile(
-                filepath=uri_to_filepath(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
-        ]
+        return [self.location_to_range_in_file(x) for x in response]
+
+    async def find_references(
+        self, position: Position, filename: str, include_declaration: bool = False
+    ) -> List[RangeInFile]:
+        response = await self.lsp_client.find_references(
+            filename,
+            position,
+            include_declaration=include_declaration,
+        )
+        return [self.location_to_range_in_file(x) for x in response["result"]]
 
     async def document_symbol(self, filepath: str) -> List:
         response = await self.lsp_client.document_symbol(filepath)
@@ -314,15 +363,55 @@ class ContinueLSPClient(BaseModel):
                 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"],
-                    ),
-                ),
+                location=self.location_to_range_in_file(x["location"]),
             )
             for x in response["result"]
         ]
+
+    async def folding_range(self, filepath: str) -> List[FoldingRange]:
+        response = await self.lsp_client.folding_range(filepath)
+
+        return [
+            FoldingRange(
+                range=Range.from_shorthand(
+                    x["startLine"],
+                    x.get("startCharacter", 0),
+                    x["endLine"] if "endCharacter" in x else x["endLine"] + 1,
+                    x.get("endCharacter", 0),
+                ),
+                kind=x.get("kind"),
+            )
+            for x in response
+        ]
+
+    async def get_enclosing_folding_range_of_position(
+        self, position: Position, filepath: str
+    ) -> Optional[FoldingRange]:
+        ranges = await self.folding_range(filepath)
+
+        max_start_position = Position(line=0, character=0)
+        max_range = None
+        for r in ranges:
+            if r.range.contains(position):
+                if r.range.start > max_start_position:
+                    max_start_position = r.range.start
+                    max_range = r
+
+        return max_range
+
+    async def get_enclosing_folding_range(
+        self, range_in_file: RangeInFile
+    ) -> Optional[FoldingRange]:
+        ranges = await self.folding_range(range_in_file.filepath)
+
+        max_start_position = Position(line=0, character=0)
+        max_range = None
+        for r in ranges:
+            if r.range.contains(range_in_file.range.start) and r.range.contains(
+                range_in_file.range.end
+            ):
+                if r.range.start > max_start_position:
+                    max_start_position = r.range.start
+                    max_range = r
+
+        return max_range