fixes #325: initialize LSP progress token before using it

Author: syphar

Diff for: pylsp/_utils.py

@@ -14,6 +14,7 @@
 import jedi
 
 JEDI_VERSION = jedi.__version__
+CALL_TIMEOUT = 10
 
 # Eol chars accepted by the LSP protocol
 # the ordering affects performance

Diff for: pylsp/python_lsp.py

@@ -152,7 +152,7 @@ class PythonLSPServer(MethodDispatcher):
 
     # pylint: disable=too-many-public-methods,redefined-builtin
 
-    def __init__(self, rx, tx, check_parent_process=False, consumer=None):
+    def __init__(self, rx, tx, check_parent_process=False, consumer=None, *, endpoint_cls=None):
         self.workspace = None
         self.config = None
         self.root_uri = None
@@ -172,11 +172,13 @@ def __init__(self, rx, tx, check_parent_process=False, consumer=None):
         else:
             self._jsonrpc_stream_writer = None
 
+        endpoint_cls = endpoint_cls or Endpoint
+
         # if consumer is None, it is assumed that the default streams-based approach is being used
         if consumer is None:
-            self._endpoint = Endpoint(self, self._jsonrpc_stream_writer.write, max_workers=MAX_WORKERS)
+            self._endpoint = endpoint_cls(self, self._jsonrpc_stream_writer.write, max_workers=MAX_WORKERS)
         else:
-            self._endpoint = Endpoint(self, consumer, max_workers=MAX_WORKERS)
+            self._endpoint = endpoint_cls(self, consumer, max_workers=MAX_WORKERS)
 
         self._dispatchers = []
         self._shutdown = False

Diff for: pylsp/workspace.py

@@ -37,6 +37,7 @@ class Workspace:
 
     M_PUBLISH_DIAGNOSTICS = 'textDocument/publishDiagnostics'
     M_PROGRESS = '$/progress'
+    M_INITIALIZE_PROGRESS = 'window/workDoneProgress/create'
     M_APPLY_EDIT = 'workspace/applyEdit'
     M_SHOW_MESSAGE = 'window/showMessage'
 
@@ -152,6 +153,9 @@ def _progress_begin(
         percentage: Optional[int] = None,
     ) -> str:
         token = str(uuid.uuid4())
+
+        self._endpoint.request(self.M_INITIALIZE_PROGRESS, {'token': token}).result(_utils.CALL_TIMEOUT)
+
         value = {
             "kind": "begin",
             "title": title,

Diff for: test/fixtures.py

@@ -5,7 +5,9 @@
 from io import StringIO
 from unittest.mock import MagicMock
 import pytest
+from pylsp_jsonrpc.dispatchers import MethodDispatcher
 from pylsp_jsonrpc.endpoint import Endpoint
+from pylsp_jsonrpc.exceptions import JsonRpcException
 
 from pylsp import uris
 from pylsp.config.config import Config
@@ -24,7 +26,7 @@ def main():
 @pytest.fixture
 def pylsp(tmpdir):
     """ Return an initialized python LS """
-    ls = PythonLSPServer(StringIO, StringIO)
+    ls = FakePythonLSPServer(StringIO, StringIO, endpoint_cls=FakeEndpoint)
 
     ls.m_initialize(
         processId=1,
@@ -38,7 +40,7 @@ def pylsp(tmpdir):
 @pytest.fixture
 def pylsp_w_workspace_folders(tmpdir):
     """ Return an initialized python LS """
-    ls = PythonLSPServer(StringIO, StringIO)
+    ls = FakePythonLSPServer(StringIO, StringIO, endpoint_cls=FakeEndpoint)
 
     folder1 = tmpdir.mkdir('folder1')
     folder2 = tmpdir.mkdir('folder2')
@@ -63,14 +65,55 @@ def pylsp_w_workspace_folders(tmpdir):
     return (ls, workspace_folders)
 
 
+class FakeEditorMethodsMixin:
+    """
+    Represents the methods to be added to a dispatcher class when faking an editor.
+    """
+    def m_window__work_done_progress__create(self, *args, **kwargs):
+        """
+        Fake editor method `window/workDoneProgress/create`.
+
+        related spec:
+        https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#window_workDoneProgress_create
+        """
+        return None
+
+
+class FakePythonLSPServer(FakeEditorMethodsMixin, PythonLSPServer):
+    pass
+
+
+class FakeEndpoint(Endpoint):
+    """
+    Fake Endpoint representing the editor / LSP client.
+
+    The `dispatcher` dict will be used to synchronously calculate the responses
+    for calls to `.request` and resolve the futures with the value or errors.
+
+    Fake methods in the `dispatcher` should raise `JsonRpcException` for any
+    error.
+    """
+    def request(self, method, params):
+        request_future = super().request(method, params)
+        try:
+            request_future.set_result(self._dispatcher[method](params))
+        except JsonRpcException as e:
+            request_future.set_exception(e)
+
+        return request_future
+
+
 @pytest.fixture()
 def consumer():
     return MagicMock()
 
 
 @pytest.fixture()
 def endpoint(consumer):  # pylint: disable=redefined-outer-name
-    return Endpoint({}, consumer, id_generator=lambda: "id")
+    class Dispatcher(FakeEditorMethodsMixin, MethodDispatcher):
+        pass
+
+    return FakeEndpoint(Dispatcher(), consumer, id_generator=lambda: "id")
 
 
 @pytest.fixture

Diff for: test/test_language_server.py

@@ -12,8 +12,8 @@
 import pytest
 
 from pylsp.python_lsp import start_io_lang_server, PythonLSPServer
+from pylsp._utils import CALL_TIMEOUT
 
-CALL_TIMEOUT = 10
 RUNNING_IN_CI = bool(os.environ.get('CI'))
 
 

Diff for: test/test_workspace.py

@@ -299,13 +299,17 @@ def test_progress_simple(workspace, consumer):
     with workspace.report_progress("some_title"):
         pass
 
+    init_call, *progress_calls = consumer.call_args_list
+
+    assert init_call[0][0]['method'] == 'window/workDoneProgress/create'
+
     # same method for all calls
-    assert all(call[0][0]["method"] == "$/progress" for call in consumer.call_args_list)
+    assert all(call[0][0]["method"] == "$/progress" for call in progress_calls), consumer.call_args_list
 
     # same token used in all calls
-    assert len({call[0][0]["params"]["token"] for call in consumer.call_args_list}) == 1
+    assert len({call[0][0]["params"]["token"] for call in progress_calls} | {init_call[0][0]['params']['token']}) == 1
 
-    assert [call[0][0]["params"]["value"] for call in consumer.call_args_list] == [
+    assert [call[0][0]["params"]["value"] for call in progress_calls] == [
         {"kind": "begin", "title": "some_title"},
         {"kind": "end"},
     ]
@@ -319,13 +323,17 @@ def test_progress_with_percent(workspace, consumer):
         progress_message("fifty", 50)
         progress_message("ninety", 90)
 
-    # same method for all calls
-    assert all(call[0][0]["method"] == "$/progress" for call in consumer.call_args_list)
+    init_call, *progress_calls = consumer.call_args_list
+
+    assert init_call[0][0]['method'] == 'window/workDoneProgress/create'
+
+    # same method for all progress calls
+    assert all(call[0][0]["method"] == "$/progress" for call in progress_calls)
 
     # same token used in all calls
-    assert len({call[0][0]["params"]["token"] for call in consumer.call_args_list}) == 1
+    assert len({call[0][0]["params"]["token"] for call in progress_calls} | {init_call[0][0]['params']['token']}) == 1
 
-    assert [call[0][0]["params"]["value"] for call in consumer.call_args_list] == [
+    assert [call[0][0]["params"]["value"] for call in progress_calls] == [
         {
             "kind": "begin",
             "message": "initial message",
@@ -353,13 +361,16 @@ class DummyError(Exception):
         # test.
         pass
 
+    init_call, *progress_calls = consumer.call_args_list
+    assert init_call[0][0]['method'] == 'window/workDoneProgress/create'
+
     # same method for all calls
-    assert all(call[0][0]["method"] == "$/progress" for call in consumer.call_args_list)
+    assert all(call[0][0]["method"] == "$/progress" for call in progress_calls)
 
     # same token used in all calls
-    assert len({call[0][0]["params"]["token"] for call in consumer.call_args_list}) == 1
+    assert len({call[0][0]["params"]["token"] for call in progress_calls} | {init_call[0][0]['params']['token']}) == 1
 
-    assert [call[0][0]["params"]["value"] for call in consumer.call_args_list] == [
+    assert [call[0][0]["params"]["value"] for call in progress_calls] == [
         {"kind": "begin", "title": "some_title"},
         {"kind": "end"},
     ]