LSP progress reporting (#456)

Author: bagel897

src/codegen/extensions/lsp/lsp.py

@@ -71,35 +71,46 @@ def did_close(server: CodegenLanguageServer, params: types.DidCloseTextDocumentP
 
 @server.feature(
     types.TEXT_DOCUMENT_RENAME,
+    options=types.RenameOptions(work_done_progress=True),
 )
 def rename(server: CodegenLanguageServer, params: types.RenameParams) -> types.RenameResult:
     symbol = server.get_symbol(params.text_document.uri, params.position)
     if symbol is None:
         logger.warning(f"No symbol found at {params.text_document.uri}:{params.position}")
         return
     logger.info(f"Renaming symbol {symbol.name} to {params.new_name}")
+    task = server.progress_manager.begin_with_token(f"Renaming symbol {symbol.name} to {params.new_name}", params.work_done_token)
     symbol.rename(params.new_name)
+    task.update("Committing changes")
     server.codebase.commit()
+    task.end()
     return server.io.get_workspace_edit()
 
 
 @server.feature(
     types.TEXT_DOCUMENT_DOCUMENT_SYMBOL,
+    options=types.DocumentSymbolOptions(work_done_progress=True),
 )
 def document_symbol(server: CodegenLanguageServer, params: types.DocumentSymbolParams) -> types.DocumentSymbolResult:
     file = server.get_file(params.text_document.uri)
     symbols = []
-    for symbol in file.symbols:
+    task = server.progress_manager.begin_with_token(f"Getting document symbols for {params.text_document.uri}", params.work_done_token, count=len(file.symbols))
+    for idx, symbol in enumerate(file.symbols):
+        task.update(f"Getting document symbols for {params.text_document.uri}", count=idx)
         symbols.append(get_document_symbol(symbol))
+    task.end()
     return symbols
 
 
 @server.feature(
     types.TEXT_DOCUMENT_DEFINITION,
+    options=types.DefinitionOptions(work_done_progress=True),
 )
 def definition(server: CodegenLanguageServer, params: types.DefinitionParams):
     node = server.get_node_under_cursor(params.text_document.uri, params.position)
+    task = server.progress_manager.begin_with_token(f"Getting definition for {params.text_document.uri}", params.work_done_token)
     resolved = go_to_definition(node, params.text_document.uri, params.position)
+    task.end()
     return types.Location(
         uri=resolved.file.path.as_uri(),
         range=get_range(resolved),
@@ -108,15 +119,11 @@ def definition(server: CodegenLanguageServer, params: types.DefinitionParams):
 
 @server.feature(
     types.TEXT_DOCUMENT_CODE_ACTION,
-    options=types.CodeActionOptions(resolve_provider=True),
+    options=types.CodeActionOptions(resolve_provider=True, work_done_progress=True),
 )
 def code_action(server: CodegenLanguageServer, params: types.CodeActionParams) -> types.CodeActionResult:
     logger.info(f"Received code action: {params}")
-    if params.context.only:
-        only = [types.CodeActionKind(kind) for kind in params.context.only]
-    else:
-        only = None
-    actions = server.get_actions_for_range(params.text_document.uri, params.range, only)
+    actions = server.get_actions_for_range(params)
     return actions
 
 

src/codegen/extensions/lsp/progress.py

@@ -0,0 +1,60 @@
+import uuid
+
+from lsprotocol import types
+from lsprotocol.types import ProgressToken
+from pygls.lsp.server import LanguageServer
+
+from codegen.sdk.codebase.progress.progress import Progress
+from codegen.sdk.codebase.progress.stub_task import StubTask
+from codegen.sdk.codebase.progress.task import Task
+
+
+class LSPTask(Task):
+    count: int | None
+
+    def __init__(self, server: LanguageServer, message: str, token: ProgressToken, count: int | None = None, create_token: bool = True) -> None:
+        self.token = token
+        if create_token:
+            server.work_done_progress.begin(self.token, types.WorkDoneProgressBegin(title=message))
+        self.server = server
+        self.message = message
+        self.count = count
+        self.create_token = create_token
+
+    def update(self, message: str, count: int | None = None) -> None:
+        if self.count is not None and count is not None:
+            percent = int(count * 100 / self.count)
+        else:
+            percent = None
+        self.server.work_done_progress.report(self.token, types.WorkDoneProgressReport(message=message, percentage=percent))
+
+    def end(self) -> None:
+        if self.create_token:
+            self.server.work_done_progress.end(self.token, value=types.WorkDoneProgressEnd())
+
+
+class LSPProgress(Progress[LSPTask | StubTask]):
+    initialized = False
+
+    def __init__(self, server: LanguageServer, initial_token: ProgressToken | None = None):
+        self.server = server
+        self.initial_token = initial_token
+        if initial_token is not None:
+            self.server.work_done_progress.begin(initial_token, types.WorkDoneProgressBegin(title="Parsing codebase..."))
+
+    def begin_with_token(self, message: str, token: ProgressToken | None = None, *, count: int | None = None, create_token: bool = True) -> LSPTask | StubTask:
+        if token is None:
+            return StubTask()
+        return LSPTask(self.server, message, token, count, create_token=create_token)
+
+    def begin(self, message: str, count: int | None = None) -> LSPTask | StubTask:
+        if self.initialized:
+            token = str(uuid.uuid4())
+            self.server.work_done_progress.create(token).result()
+            return LSPTask(self.server, message, token, count, create_token=False)
+        return self.begin_with_token(message, self.initial_token, count=None, create_token=False)
+
+    def finish_initialization(self) -> None:
+        self.initialized = False  # We can't initiate server work during syncs
+        if self.initial_token is not None:
+            self.server.work_done_progress.end(self.initial_token, value=types.WorkDoneProgressEnd())

src/codegen/extensions/lsp/protocol.py

@@ -2,10 +2,11 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from lsprotocol.types import INITIALIZE, InitializeParams, InitializeResult, WorkDoneProgressBegin, WorkDoneProgressEnd
+from lsprotocol.types import INITIALIZE, InitializeParams, InitializeResult
 from pygls.protocol import LanguageServerProtocol, lsp_method
 
 from codegen.extensions.lsp.io import LSPIO
+from codegen.extensions.lsp.progress import LSPProgress
 from codegen.extensions.lsp.utils import get_path
 from codegen.sdk.codebase.config import CodebaseConfig
 from codegen.sdk.core.codebase import Codebase
@@ -19,6 +20,7 @@ class CodegenLanguageServerProtocol(LanguageServerProtocol):
     _server: "CodegenLanguageServer"
 
     def _init_codebase(self, params: InitializeParams) -> None:
+        progress = LSPProgress(self._server, params.work_done_token)
         if params.root_path:
             root = Path(params.root_path)
         elif params.root_uri:
@@ -27,15 +29,13 @@ def _init_codebase(self, params: InitializeParams) -> None:
             root = os.getcwd()
         config = CodebaseConfig(feature_flags=CodebaseFeatureFlags(full_range_index=True))
         io = LSPIO(self.workspace)
-        self._server.codebase = Codebase(repo_path=str(root), config=config, io=io)
+        self._server.codebase = Codebase(repo_path=str(root), config=config, io=io, progress=progress)
+        self._server.progress_manager = progress
         self._server.io = io
-        if params.work_done_token:
-            self._server.work_done_progress.end(params.work_done_token, WorkDoneProgressEnd(message="Parsing codebase..."))
+        progress.finish_initialization()
 
     @lsp_method(INITIALIZE)
     def lsp_initialize(self, params: InitializeParams) -> InitializeResult:
         ret = super().lsp_initialize(params)
-        if params.work_done_token:
-            self._server.work_done_progress.begin(params.work_done_token, WorkDoneProgressBegin(title="Parsing codebase..."))
         self._init_codebase(params)
         return ret

src/codegen/extensions/lsp/server.py

@@ -1,5 +1,4 @@
 import logging
-from collections.abc import Sequence
 from typing import Any, Optional
 
 from lsprotocol import types
@@ -8,8 +7,9 @@
 
 from codegen.extensions.lsp.codemods import ACTIONS
 from codegen.extensions.lsp.codemods.base import CodeAction
-from codegen.extensions.lsp.execute import execute_action, get_execute_action
+from codegen.extensions.lsp.execute import execute_action
 from codegen.extensions.lsp.io import LSPIO
+from codegen.extensions.lsp.progress import LSPProgress
 from codegen.extensions.lsp.range import get_tree_sitter_range
 from codegen.extensions.lsp.utils import get_path
 from codegen.sdk.core.codebase import Codebase
@@ -23,13 +23,14 @@
 class CodegenLanguageServer(LanguageServer):
     codebase: Optional[Codebase]
     io: Optional[LSPIO]
+    progress_manager: Optional[LSPProgress]
     actions: dict[str, CodeAction]
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self.actions = {action.command_name(): action for action in ACTIONS}
-        for action in self.actions.values():
-            self.command(action.command_name())(get_execute_action(action))
+        # for action in self.actions.values():
+        #     self.command(action.command_name())(get_execute_action(action))
 
     def get_file(self, uri: str) -> SourceFile | File:
         path = get_path(uri)
@@ -68,19 +69,25 @@ def get_node_for_range(self, uri: str, range: Range) -> Editable | None:
             return node
         return None
 
-    def get_actions_for_range(self, uri: str, range: Range, only: Sequence[types.CodeActionKind] | None = None) -> list[types.CodeAction]:
-        node = self.get_node_under_cursor(uri, range.start, range.end)
+    def get_actions_for_range(self, params: types.CodeActionParams) -> list[types.CodeAction]:
+        if params.context.only is not None:
+            only = [types.CodeActionKind(kind) for kind in params.context.only]
+        else:
+            only = None
+        node = self.get_node_under_cursor(params.text_document.uri, params.range.start)
         if node is None:
-            logger.warning(f"No node found for range {range} in {uri}")
+            logger.warning(f"No node found for range {params.range} in {params.text_document.uri}")
             return []
         actions = []
-        for action in self.actions.values():
+        task = self.progress_manager.begin_with_token(f"Getting code actions for {params.text_document.uri}", params.work_done_token, count=len(self.actions))
+        for idx, action in enumerate(self.actions.values()):
+            task.update(f"Checking action {action.name}", idx)
             if only and action.kind not in only:
                 logger.warning(f"Skipping action {action.kind} because it is not in {only}")
                 continue
             if action.is_applicable(self, node):
-                actions.append(action.to_lsp(uri, range))
-
+                actions.append(action.to_lsp(params.text_document.uri, params.range))
+        task.end()
         return actions
 
     def resolve_action(self, action: types.CodeAction) -> types.CodeAction:

src/codegen/sdk/codebase/codebase_context.py

@@ -16,6 +16,7 @@
 from codegen.sdk.codebase.diff_lite import ChangeType, DiffLite
 from codegen.sdk.codebase.flagging.flags import Flags
 from codegen.sdk.codebase.io.file_io import FileIO
+from codegen.sdk.codebase.progress.stub_progress import StubProgress
 from codegen.sdk.codebase.transaction_manager import TransactionManager
 from codegen.sdk.codebase.validation import get_edges, post_reset_validation
 from codegen.sdk.core.autocommit import AutoCommit, commiter
@@ -39,6 +40,7 @@
     from codegen.git.repo_operator.repo_operator import RepoOperator
     from codegen.sdk.codebase.io.io import IO
     from codegen.sdk.codebase.node_classes.node_classes import NodeClasses
+    from codegen.sdk.codebase.progress.progress import Progress
     from codegen.sdk.core.dataclasses.usage import Usage
     from codegen.sdk.core.expressions import Expression
     from codegen.sdk.core.external_module import ExternalModule
@@ -111,16 +113,19 @@ class CodebaseContext:
     projects: list[ProjectConfig]
     unapplied_diffs: list[DiffLite]
     io: IO
+    progress: Progress
 
     def __init__(
         self,
         projects: list[ProjectConfig],
         config: CodebaseConfig = DefaultConfig,
         io: IO | None = None,
+        progress: Progress | None = None,
     ) -> None:
         """Initializes codebase graph and TransactionManager"""
         from codegen.sdk.core.parser import Parser
 
+        self.progress = progress or StubProgress()
         self._graph = PyDiGraph()
         self.filepath_idx = {}
         self._ext_module_idx = {}
@@ -371,7 +376,6 @@ def _process_diff_files(self, files_to_sync: Mapping[SyncType, list[Path]], incr
         skip_uncache = incremental and ((len(files_to_sync[SyncType.DELETE]) + len(files_to_sync[SyncType.REPARSE])) == 0)
         if not skip_uncache:
             uncache_all()
-
         # Step 0: Start the dependency manager and language engine if they exist
         # Start the dependency manager. This may or may not run asynchronously, depending on the implementation
         if self.dependency_manager is not None:
@@ -429,24 +433,29 @@ def _process_diff_files(self, files_to_sync: Mapping[SyncType, list[Path]], incr
             file = self.get_file(file_path)
             file.remove_internal_edges()
 
+        task = self.progress.begin("Reparsing updated files", count=len(files_to_sync[SyncType.REPARSE]))
         files_to_resolve = []
         # Step 4: Reparse updated files
-        for file_path in files_to_sync[SyncType.REPARSE]:
+        for idx, file_path in enumerate(files_to_sync[SyncType.REPARSE]):
+            task.update(f"Reparsing {self.to_relative(file_path)}", count=idx)
             file = self.get_file(file_path)
             to_resolve.extend(file.unparse(reparse=True))
             to_resolve = list(filter(lambda node: self.has_node(node.node_id) and node is not None, to_resolve))
             file.sync_with_file_content()
             files_to_resolve.append(file)
-
+        task.end()
         # Step 5: Add new files as nodes to graph (does not yet add edges)
-        for filepath in files_to_sync[SyncType.ADD]:
+        task = self.progress.begin("Adding new files", count=len(files_to_sync[SyncType.ADD]))
+        for idx, filepath in enumerate(files_to_sync[SyncType.ADD]):
+            task.update(f"Adding {self.to_relative(filepath)}", count=idx)
             content = self.io.read_text(filepath)
             # TODO: this is wrong with context changes
             if filepath.suffix in self.extensions:
                 file_cls = self.node_classes.file_cls
                 new_file = file_cls.from_content(filepath, content, self, sync=False, verify_syntax=False)
                 if new_file is not None:
                     files_to_resolve.append(new_file)
+        task.end()
         for file in files_to_resolve:
             to_resolve.append(file)
             to_resolve.extend(file.get_nodes())
@@ -474,27 +483,35 @@ def _process_diff_files(self, files_to_sync: Mapping[SyncType, list[Path]], incr
             self._computing = True
             try:
                 logger.info(f"> Computing import resolution edges for {counter[NodeType.IMPORT]} imports")
+                task = self.progress.begin("Resolving imports", count=counter[NodeType.IMPORT])
                 for node in to_resolve:
                     if node.node_type == NodeType.IMPORT:
+                        task.update(f"Resolving imports in {node.filepath}", count=idx)
                         node._remove_internal_edges(EdgeType.IMPORT_SYMBOL_RESOLUTION)
                         node.add_symbol_resolution_edge()
                         to_resolve.extend(node.symbol_usages)
+                task.end()
                 if counter[NodeType.EXPORT] > 0:
                     logger.info(f"> Computing export dependencies for {counter[NodeType.EXPORT]} exports")
+                    task = self.progress.begin("Computing export dependencies", count=counter[NodeType.EXPORT])
                     for node in to_resolve:
                         if node.node_type == NodeType.EXPORT:
+                            task.update(f"Computing export dependencies for {node.filepath}", count=idx)
                             node._remove_internal_edges(EdgeType.EXPORT)
                             node.compute_export_dependencies()
                             to_resolve.extend(node.symbol_usages)
+                    task.end()
                 if counter[NodeType.SYMBOL] > 0:
                     from codegen.sdk.core.interfaces.inherits import Inherits
 
                     logger.info("> Computing superclass dependencies")
+                    task = self.progress.begin("Computing superclass dependencies", count=counter[NodeType.SYMBOL])
                     for symbol in to_resolve:
                         if isinstance(symbol, Inherits):
+                            task.update(f"Computing superclass dependencies for {symbol.filepath}", count=idx)
                             symbol._remove_internal_edges(EdgeType.SUBCLASS)
                             symbol.compute_superclass_dependencies()
-
+                    task.end()
                 if not skip_uncache:
                     uncache_all()
                 self._compute_dependencies(to_resolve, incremental)
@@ -504,17 +521,20 @@ def _process_diff_files(self, files_to_sync: Mapping[SyncType, list[Path]], incr
     def _compute_dependencies(self, to_update: list[Importable], incremental: bool):
         seen = set()
         while to_update:
+            task = self.progress.begin("Computing dependencies", count=len(to_update))
             step = to_update.copy()
             to_update.clear()
             logger.info(f"> Incrementally computing dependencies for {len(step)} nodes")
-            for current in step:
+            for idx, current in enumerate(step):
+                task.update(f"Computing dependencies for {current.filepath}", count=idx)
                 if current not in seen:
                     seen.add(current)
                     to_update.extend(current.recompute(incremental))
             if not incremental:
                 for node in self._graph.nodes():
                     if node not in seen:
                         to_update.append(node)
+            task.end()
         seen.clear()
 
     def build_subgraph(self, nodes: list[NodeId]) -> PyDiGraph[Importable, Edge]:

src/codegen/sdk/codebase/progress/progress.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+if TYPE_CHECKING:
+    from codegen.sdk.codebase.progress.task import Task
+
+T = TypeVar("T", bound="Task")
+
+
+class Progress(ABC, Generic[T]):
+    @abstractmethod
+    def begin(self, message: str, count: int | None = None) -> T:
+        pass

src/codegen/sdk/codebase/progress/stub_progress.py

@@ -0,0 +1,7 @@
+from codegen.sdk.codebase.progress.progress import Progress
+from codegen.sdk.codebase.progress.stub_task import StubTask
+
+
+class StubProgress(Progress[StubTask]):
+    def begin(self, message: str, count: int | None = None) -> StubTask:
+        return StubTask()