Workaround for pygls performance issue

extensions/positron-python/python_files/posit/positron/positron_jedilsp.py

@@ -38,6 +38,7 @@
     workspace_symbol,
 )
 from ._vendor.lsprotocol.types import (
+    CANCEL_REQUEST,
     COMPLETION_ITEM_RESOLVE,
     INITIALIZE,
     TEXT_DOCUMENT_CODE_ACTION,
@@ -167,6 +168,12 @@ class PositronInitializationOptions:
 
 
 class PositronJediLanguageServerProtocol(JediLanguageServerProtocol):
+    def __init__(self, server, converter):
+        super().__init__(server, converter)
+
+        # See `self._data_received` for a description.
+        self._messages_to_handle = []
+
     @lru_cache  # noqa: B019
     def get_message_type(self, method: str) -> Optional[Type]:
         # Overriden to include custom Positron LSP messages.
@@ -216,6 +223,52 @@ def lsp_initialize(self, params: InitializeParams) -> InitializeResult:
 
         return result
 
+    def _data_received(self, data: bytes) -> None:
+        # Workaround to a pygls performance issue where the server still executes requests
+        # even if they're immediately cancelled.
+        # See: https://github.com/openlawlibrary/pygls/issues/517.
+
+        # This should parse `data` and call `self._procedure_handler` with each parsed message.
+        # That usually handles each message, but we've overridden it to just add them to a queue.
+        self._messages_to_handle = []
+        super()._data_received(data)
+
+        def is_request(message):
+            return hasattr(message, "method") and hasattr(message, "id")
+
+        def is_cancel_notification(message):
+            return getattr(message, "method", None) == CANCEL_REQUEST
+
+        # First pass: find all requests that were cancelled in the same batch of `data`.
+        request_ids = set()
+        cancelled_ids = set()
+        for message in self._messages_to_handle:
+            if is_request(message):
+                request_ids.add(message.id)
+            elif is_cancel_notification(message) and message.params.id in request_ids:
+                cancelled_ids.add(message.params.id)
+
+        # Second pass: remove all requests that were cancelled in the same batch of `data`,
+        # and the cancel notifications themselves.
+        self._messages_to_handle = [
+            msg
+            for msg in self._messages_to_handle
+            if not (
+                # Remove cancel notifications whose params.id is in cancelled_ids...
+                (is_cancel_notification(msg) and msg.params.id in cancelled_ids)
+                # ...or original messages whose id is in cancelled_ids.
+                or (is_request(msg) and msg.id in cancelled_ids)
+            )
+        ]
+
+        # Now handle the messages.
+        for message in self._messages_to_handle:
+            super()._procedure_handler(message)
+
+    def _procedure_handler(self, message) -> None:
+        # Overridden to just queue messages which are handled later in `self._data_received`.
+        self._messages_to_handle.append(message)
+
 
 class PositronJediLanguageServer(JediLanguageServer):
     """Positron extension to the Jedi language server."""
@@ -636,6 +689,7 @@ def positron_highlight(
     return highlight(server, params)
 
 
+@POSITRON.thread()
 @POSITRON.feature(TEXT_DOCUMENT_HOVER)
 def positron_hover(
     server: PositronJediLanguageServer, params: TextDocumentPositionParams

extensions/positron-python/python_files/posit/positron/tests/test_positron_jedilsp.py

@@ -3,6 +3,7 @@
 # Licensed under the Elastic License 2.0. See LICENSE.txt for license information.
 #
 
+import json
 import os
 from functools import partial
 from pathlib import Path
@@ -16,6 +17,9 @@
 from positron._vendor import cattrs
 from positron._vendor.jedi_language_server import jedi_utils
 from positron._vendor.lsprotocol.types import (
+    TEXT_DOCUMENT_HOVER,
+    CancelParams,
+    CancelRequestNotification,
     ClientCapabilities,
     CompletionClientCapabilities,
     CompletionClientCapabilitiesCompletionItemType,
@@ -26,6 +30,7 @@
     DocumentSymbol,
     DocumentSymbolParams,
     Hover,
+    HoverParams,
     InitializeParams,
     Location,
     MarkupContent,
@@ -40,6 +45,7 @@
     SymbolKind,
     TextDocumentClientCapabilities,
     TextDocumentEdit,
+    TextDocumentHoverRequest,
     TextDocumentIdentifier,
     TextDocumentItem,
     TextDocumentPositionParams,
@@ -1042,3 +1048,37 @@ def test_positron_rename(
     text_document_edit = workspace_edit.document_changes[0]
     assert isinstance(text_document_edit, TextDocumentEdit)
     assert text_document_edit.edits == expected_text_edits
+
+
+def test_cancel_request_immediately() -> None:
+    # Test our workaround to a pygls performance issue where the server still executes requests
+    # even if they're immediately cancelled.
+    # See: https://github.com/openlawlibrary/pygls/issues/517.
+    server = create_server()
+    protocol = server.lsp
+
+    # Register a textDocument/hover handler and track whether it executes.
+    did_hover_executed = False
+
+    @server.feature(TEXT_DOCUMENT_HOVER)
+    def _hover(_server: PositronJediLanguageServer, _params: TextDocumentPositionParams):
+        nonlocal did_hover_executed
+        did_hover_executed = True
+
+    # Helper to encode LSP messages, adapted from `pygls.json_rpc.JsonRPCProtocol._send_data`.
+    def encode(message):
+        body = json.dumps(message, default=protocol._serialize_message)  # noqa: SLF001
+        header = f"Content-Length: {len(body)}\r\n\r\n"
+        return (header + body).encode(protocol.CHARSET)
+
+    # Send a hover request and immediately cancel it in the same payload.
+    hover_request = TextDocumentHoverRequest(
+        0, HoverParams(TextDocumentIdentifier(""), Position(0, 0))
+    )
+    cancel_request = CancelRequestNotification(CancelParams(hover_request.id))
+    data = encode(hover_request) + encode(cancel_request)
+
+    # Call `_data_received` instead of `data_received` so that errors are raised.
+    protocol._data_received(data)  # noqa: SLF001
+
+    assert not did_hover_executed