Add more progress support

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

@@ -42,20 +42,19 @@ def __init__(self, server: LanguageServer, initial_token: ProgressToken | None =
         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, *, count: int | None = None) -> LSPTask:
+    def begin_with_token(self, message: str, token: ProgressToken | None = None, *, count: int | None = None) -> LSPTask | StubTask:
+        if token is None:
+            return StubTask()
         return LSPTask(self.server, message, token, count, create_token=False)
 
     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)
-        elif self.initial_token is not None:
-            return self.begin_with_token(message, self.initial_token, count=None)
-        else:
-            return StubTask()
+        return self.begin_with_token(message, self.initial_token, count=None)
 
     def finish_initialization(self) -> None:
-        self.initialized = True
+        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

@@ -30,6 +30,7 @@ def _init_codebase(self, params: InitializeParams) -> None:
         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, progress=progress)
+        self._server.progress_manager = progress
         self._server.io = io
         progress.finish_initialization()
 

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: