jaclang 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

jaclang/vendor/pygls/protocol.py

@@ -0,0 +1,1184 @@
+############################################################################
+# Copyright(c) Open Law Library. All rights reserved.                      #
+# See ThirdPartyNotices.txt in the project root for additional notices.    #
+#                                                                          #
+# Licensed under the Apache License, Version 2.0 (the "License")           #
+# you may not use this file except in compliance with the License.         #
+# You may obtain a copy of the License at                                  #
+#                                                                          #
+#     http: // www.apache.org/licenses/LICENSE-2.0                         #
+#                                                                          #
+# Unless required by applicable law or agreed to in writing, software      #
+# distributed under the License is distributed on an "AS IS" BASIS,        #
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. #
+# See the License for the specific language governing permissions and      #
+# limitations under the License.                                           #
+############################################################################
+from __future__ import annotations
+import asyncio
+import enum
+import functools
+import json
+import logging
+import re
+import sys
+import uuid
+import traceback
+from collections import namedtuple
+from concurrent.futures import Future
+from functools import lru_cache, partial
+from itertools import zip_longest
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    TYPE_CHECKING,
+)
+
+if TYPE_CHECKING:
+    from jaclang.vendor.pygls.server import LanguageServer, WebSocketTransportAdapter
+
+
+import attrs
+from cattrs.errors import ClassValidationError
+from jaclang.vendor.pygls.lsprotocol import converters
+
+from jaclang.vendor.pygls.capabilities import ServerCapabilitiesBuilder
+from jaclang.vendor.pygls.constants import ATTR_FEATURE_TYPE
+from jaclang.vendor.pygls.exceptions import (
+    JsonRpcException,
+    JsonRpcInternalError,
+    JsonRpcInvalidParams,
+    JsonRpcMethodNotFound,
+    JsonRpcRequestCancelled,
+    FeatureNotificationError,
+    FeatureRequestError,
+)
+from jaclang.vendor.pygls.feature_manager import (
+    FeatureManager,
+    assign_help_attrs,
+    is_thread_function,
+)
+from jaclang.vendor.pygls.lsp import ConfigCallbackType, ShowDocumentCallbackType
+from jaclang.vendor.pygls.lsprotocol.types import (
+    CANCEL_REQUEST,
+    CLIENT_REGISTER_CAPABILITY,
+    CLIENT_UNREGISTER_CAPABILITY,
+    EXIT,
+    INITIALIZE,
+    INITIALIZED,
+    METHOD_TO_TYPES,
+    NOTEBOOK_DOCUMENT_DID_CHANGE,
+    NOTEBOOK_DOCUMENT_DID_CLOSE,
+    NOTEBOOK_DOCUMENT_DID_OPEN,
+    LOG_TRACE,
+    SET_TRACE,
+    SHUTDOWN,
+    TEXT_DOCUMENT_DID_CHANGE,
+    TEXT_DOCUMENT_DID_CLOSE,
+    TEXT_DOCUMENT_DID_OPEN,
+    TEXT_DOCUMENT_PUBLISH_DIAGNOSTICS,
+    WINDOW_LOG_MESSAGE,
+    WINDOW_SHOW_DOCUMENT,
+    WINDOW_SHOW_MESSAGE,
+    WINDOW_WORK_DONE_PROGRESS_CANCEL,
+    WORKSPACE_APPLY_EDIT,
+    WORKSPACE_CONFIGURATION,
+    WORKSPACE_DID_CHANGE_WORKSPACE_FOLDERS,
+    WORKSPACE_EXECUTE_COMMAND,
+    WORKSPACE_SEMANTIC_TOKENS_REFRESH,
+)
+from jaclang.vendor.pygls.lsprotocol.types import (
+    ApplyWorkspaceEditParams,
+    Diagnostic,
+    DidChangeNotebookDocumentParams,
+    DidChangeTextDocumentParams,
+    DidChangeWorkspaceFoldersParams,
+    DidCloseNotebookDocumentParams,
+    DidCloseTextDocumentParams,
+    DidOpenNotebookDocumentParams,
+    DidOpenTextDocumentParams,
+    ExecuteCommandParams,
+    InitializeParams,
+    InitializeResult,
+    LogMessageParams,
+    LogTraceParams,
+    MessageType,
+    PublishDiagnosticsParams,
+    RegistrationParams,
+    ResponseErrorMessage,
+    SetTraceParams,
+    ShowDocumentParams,
+    ShowMessageParams,
+    TraceValues,
+    UnregistrationParams,
+    WorkspaceApplyEditResponse,
+    WorkspaceEdit,
+    InitializeResultServerInfoType,
+    WorkspaceConfigurationParams,
+    WorkDoneProgressCancelParams,
+)
+from jaclang.vendor.pygls.uris import from_fs_path
+from jaclang.vendor.pygls.workspace import Workspace
+
+logger = logging.getLogger(__name__)
+
+F = TypeVar("F", bound=Callable)
+
+
+def call_user_feature(base_func, method_name):
+    """Wraps generic LSP features and calls user registered feature
+    immediately after it.
+    """
+
+    @functools.wraps(base_func)
+    def decorator(self, *args, **kwargs):
+        ret_val = base_func(self, *args, **kwargs)
+
+        try:
+            user_func = self.fm.features[method_name]
+            self._execute_notification(user_func, *args, **kwargs)
+        except KeyError:
+            pass
+        except Exception:
+            logger.exception(
+                'Failed to handle user defined notification "%s": %s', method_name, args
+            )
+
+        return ret_val
+
+    return decorator
+
+
+@attrs.define
+class JsonRPCNotification:
+    """A class that represents a generic json rpc notification message.
+    Used as a fallback for unknown types.
+    """
+
+    method: str
+    jsonrpc: str
+    params: Any
+
+
+@attrs.define
+class JsonRPCRequestMessage:
+    """A class that represents a generic json rpc request message.
+    Used as a fallback for unknown types.
+    """
+
+    id: Union[int, str]
+    method: str
+    jsonrpc: str
+    params: Any
+
+
+@attrs.define
+class JsonRPCResponseMessage:
+    """A class that represents a generic json rpc response message.
+    Used as a fallback for unknown types.
+    """
+
+    id: Union[int, str]
+    jsonrpc: str
+    result: Any
+
+
+def _dict_to_object(d: Any):
+    """Create nested objects (namedtuple) from dict."""
+
+    if d is None:
+        return None
+
+    if not isinstance(d, dict):
+        return d
+
+    type_name = d.pop("type_name", "Object")
+    return json.loads(
+        json.dumps(d),
+        object_hook=lambda p: namedtuple(type_name, p.keys(), rename=True)(*p.values()),
+    )
+
+
+def _params_field_structure_hook(obj, cls):
+    if "params" in obj:
+        obj["params"] = _dict_to_object(obj["params"])
+
+    return cls(**obj)
+
+
+def _result_field_structure_hook(obj, cls):
+    if "result" in obj:
+        obj["result"] = _dict_to_object(obj["result"])
+
+    return cls(**obj)
+
+
+def default_converter():
+    """Default converter factory function."""
+
+    converter = converters.get_converter()
+    converter.register_structure_hook(
+        JsonRPCRequestMessage, _params_field_structure_hook
+    )
+
+    converter.register_structure_hook(
+        JsonRPCResponseMessage, _result_field_structure_hook
+    )
+
+    converter.register_structure_hook(JsonRPCNotification, _params_field_structure_hook)
+
+    return converter
+
+
+class JsonRPCProtocol(asyncio.Protocol):
+    """Json RPC protocol implementation using on top of `asyncio.Protocol`.
+
+    Specification of the protocol can be found here:
+        https://www.jsonrpc.org/specification
+
+    This class provides bidirectional communication which is needed for LSP.
+    """
+
+    CHARSET = "utf-8"
+    CONTENT_TYPE = "application/vscode-jsonrpc"
+
+    MESSAGE_PATTERN = re.compile(
+        rb"^(?:[^\r\n]+\r\n)*"
+        + rb"Content-Length: (?P<length>\d+)\r\n"
+        + rb"(?:[^\r\n]+\r\n)*\r\n"
+        + rb"(?P<body>{.*)",
+        re.DOTALL,
+    )
+
+    VERSION = "2.0"
+
+    def __init__(self, server: LanguageServer, converter):
+        self._server = server
+        self._converter = converter
+
+        self._shutdown = False
+
+        # Book keeping for in-flight requests
+        self._request_futures: Dict[str, Future[Any]] = {}
+        self._result_types: Dict[str, Any] = {}
+
+        self.fm = FeatureManager(server)
+        self.transport: Optional[
+            Union[asyncio.WriteTransport, WebSocketTransportAdapter]
+        ] = None
+        self._message_buf: List[bytes] = []
+
+        self._send_only_body = False
+
+    def __call__(self):
+        return self
+
+    def _execute_notification(self, handler, *params):
+        """Executes notification message handler."""
+        if asyncio.iscoroutinefunction(handler):
+            future = asyncio.ensure_future(handler(*params))
+            future.add_done_callback(self._execute_notification_callback)
+        else:
+            if is_thread_function(handler):
+                self._server.thread_pool.apply_async(handler, (*params,))
+            else:
+                handler(*params)
+
+    def _execute_notification_callback(self, future):
+        """Success callback used for coroutine notification message."""
+        if future.exception():
+            try:
+                raise future.exception()
+            except Exception:
+                error = JsonRpcInternalError.of(sys.exc_info()).to_dict()
+                logger.exception('Exception occurred in notification: "%s"', error)
+
+            # Revisit. Client does not support response with msg_id = None
+            # https://stackoverflow.com/questions/31091376/json-rpc-2-0-allow-notifications-to-have-an-error-response
+            # self._send_response(None, error=error)
+
+    def _execute_request(self, msg_id, handler, params):
+        """Executes request message handler."""
+
+        if asyncio.iscoroutinefunction(handler):
+            future = asyncio.ensure_future(handler(params))
+            self._request_futures[msg_id] = future
+            future.add_done_callback(partial(self._execute_request_callback, msg_id))
+        else:
+            # Can't be canceled
+            if is_thread_function(handler):
+                self._server.thread_pool.apply_async(
+                    handler,
+                    (params,),
+                    callback=partial(
+                        self._send_response,
+                        msg_id,
+                    ),
+                    error_callback=partial(self._execute_request_err_callback, msg_id),
+                )
+            else:
+                self._send_response(msg_id, handler(params))
+
+    def _execute_request_callback(self, msg_id, future):
+        """Success callback used for coroutine request message."""
+        try:
+            if not future.cancelled():
+                self._send_response(msg_id, result=future.result())
+            else:
+                self._send_response(
+                    msg_id,
+                    error=JsonRpcRequestCancelled(
+                        f'Request with id "{msg_id}" is canceled'
+                    ),
+                )
+            self._request_futures.pop(msg_id, None)
+        except Exception:
+            error = JsonRpcInternalError.of(sys.exc_info()).to_dict()
+            logger.exception('Exception occurred for message "%s": %s', msg_id, error)
+            self._send_response(msg_id, error=error)
+
+    def _execute_request_err_callback(self, msg_id, exc):
+        """Error callback used for coroutine request message."""
+        exc_info = (type(exc), exc, None)
+        error = JsonRpcInternalError.of(exc_info).to_dict()
+        logger.exception('Exception occurred for message "%s": %s', msg_id, error)
+        self._send_response(msg_id, error=error)
+
+    def _get_handler(self, feature_name):
+        """Returns builtin or used defined feature by name if exists."""
+        try:
+            return self.fm.builtin_features[feature_name]
+        except KeyError:
+            try:
+                return self.fm.features[feature_name]
+            except KeyError:
+                raise JsonRpcMethodNotFound.of(feature_name)
+
+    def _handle_cancel_notification(self, msg_id):
+        """Handles a cancel notification from the client."""
+        future = self._request_futures.pop(msg_id, None)
+
+        if not future:
+            logger.warning('Cancel notification for unknown message id "%s"', msg_id)
+            return
+
+        # Will only work if the request hasn't started executing
+        if future.cancel():
+            logger.info('Cancelled request with id "%s"', msg_id)
+
+    def _handle_notification(self, method_name, params):
+        """Handles a notification from the client."""
+        if method_name == CANCEL_REQUEST:
+            self._handle_cancel_notification(params.id)
+            return
+
+        try:
+            handler = self._get_handler(method_name)
+            self._execute_notification(handler, params)
+        except (KeyError, JsonRpcMethodNotFound):
+            logger.warning('Ignoring notification for unknown method "%s"', method_name)
+        except Exception as error:
+            logger.exception(
+                'Failed to handle notification "%s": %s',
+                method_name,
+                params,
+                exc_info=True,
+            )
+            self._server._report_server_error(error, FeatureNotificationError)
+
+    def _handle_request(self, msg_id, method_name, params):
+        """Handles a request from the client."""
+        try:
+            handler = self._get_handler(method_name)
+
+            # workspace/executeCommand is a special case
+            if method_name == WORKSPACE_EXECUTE_COMMAND:
+                handler(params, msg_id)
+            else:
+                self._execute_request(msg_id, handler, params)
+
+        except JsonRpcException as error:
+            logger.exception(
+                "Failed to handle request %s %s %s",
+                msg_id,
+                method_name,
+                params,
+                exc_info=True,
+            )
+            self._send_response(msg_id, None, error.to_dict())
+            self._server._report_server_error(error, FeatureRequestError)
+        except Exception as error:
+            logger.exception(
+                "Failed to handle request %s %s %s",
+                msg_id,
+                method_name,
+                params,
+                exc_info=True,
+            )
+            err = JsonRpcInternalError.of(sys.exc_info()).to_dict()
+            self._send_response(msg_id, None, err)
+            self._server._report_server_error(error, FeatureRequestError)
+
+    def _handle_response(self, msg_id, result=None, error=None):
+        """Handles a response from the client."""
+        future = self._request_futures.pop(msg_id, None)
+
+        if not future:
+            logger.warning('Received response to unknown message id "%s"', msg_id)
+            return
+
+        if error is not None:
+            logger.debug('Received error response to message "%s": %s', msg_id, error)
+            future.set_exception(JsonRpcException.from_error(error))
+        else:
+            logger.debug('Received result for message "%s": %s', msg_id, result)
+            future.set_result(result)
+
+    def _serialize_message(self, data):
+        """Function used to serialize data sent to the client."""
+
+        if hasattr(data, "__attrs_attrs__"):
+            return self._converter.unstructure(data)
+
+        if isinstance(data, enum.Enum):
+            return data.value
+
+        return data.__dict__
+
+    def _deserialize_message(self, data):
+        """Function used to deserialize data recevied from the client."""
+
+        if "jsonrpc" not in data:
+            return data
+
+        try:
+            if "id" in data:
+                if "error" in data:
+                    return self._converter.structure(data, ResponseErrorMessage)
+                elif "method" in data:
+                    request_type = (
+                        self.get_message_type(data["method"]) or JsonRPCRequestMessage
+                    )
+                    return self._converter.structure(data, request_type)
+                else:
+                    response_type = (
+                        self._result_types.pop(data["id"]) or JsonRPCResponseMessage
+                    )
+                    return self._converter.structure(data, response_type)
+
+            else:
+                method = data.get("method", "")
+                notification_type = self.get_message_type(method) or JsonRPCNotification
+                return self._converter.structure(data, notification_type)
+
+        except ClassValidationError as exc:
+            logger.error("Unable to deserialize message\n%s", traceback.format_exc())
+            raise JsonRpcInvalidParams() from exc
+
+        except Exception as exc:
+            logger.error("Unable to deserialize message\n%s", traceback.format_exc())
+            raise JsonRpcInternalError() from exc
+
+    def _procedure_handler(self, message):
+        """Delegates message to handlers depending on message type."""
+
+        if message.jsonrpc != JsonRPCProtocol.VERSION:
+            logger.warning('Unknown message "%s"', message)
+            return
+
+        if self._shutdown and getattr(message, "method", "") != EXIT:
+            logger.warning("Server shutting down. No more requests!")
+            return
+
+        if hasattr(message, "method"):
+            if hasattr(message, "id"):
+                logger.debug("Request message received.")
+                self._handle_request(message.id, message.method, message.params)
+            else:
+                logger.debug("Notification message received.")
+                self._handle_notification(message.method, message.params)
+        else:
+            if hasattr(message, "error"):
+                logger.debug("Error message received.")
+                self._handle_response(message.id, None, message.error)
+            else:
+                logger.debug("Response message received.")
+                self._handle_response(message.id, message.result)
+
+    def _send_data(self, data):
+        """Sends data to the client."""
+        if not data:
+            return
+
+        if self.transport is None:
+            logger.error("Unable to send data, no available transport!")
+            return
+
+        try:
+            body = json.dumps(data, default=self._serialize_message)
+            logger.info("Sending data: %s", body)
+
+            if self._send_only_body:
+                # Mypy/Pyright seem to think `write()` wants `"bytes | bytearray | memoryview"`
+                # But runtime errors with anything but `str`.
+                self.transport.write(body)  # type: ignore
+                return
+
+            header = (
+                f"Content-Length: {len(body)}\r\n"
+                f"Content-Type: {self.CONTENT_TYPE}; charset={self.CHARSET}\r\n\r\n"
+            ).encode(self.CHARSET)
+
+            self.transport.write(header + body.encode(self.CHARSET))
+        except Exception as error:
+            logger.exception("Error sending data", exc_info=True)
+            self._server._report_server_error(error, JsonRpcInternalError)
+
+    def _send_response(self, msg_id, result=None, error=None):
+        """Sends a JSON RPC response to the client.
+
+        Args:
+            msg_id(str): Id from request
+            result(any): Result returned by handler
+            error(any): Error returned by handler
+        """
+
+        if error is not None:
+            response = ResponseErrorMessage(id=msg_id, error=error)
+
+        else:
+            response_type = self._result_types.pop(msg_id, JsonRPCResponseMessage)
+            response = response_type(
+                id=msg_id, result=result, jsonrpc=JsonRPCProtocol.VERSION
+            )
+
+        self._send_data(response)
+
+    def connection_lost(self, exc):
+        """Method from base class, called when connection is lost, in which case we
+        want to shutdown the server's process as well.
+        """
+        logger.error("Connection to the client is lost! Shutting down the server.")
+        sys.exit(1)
+
+    def connection_made(  # type: ignore # see: https://github.com/python/typeshed/issues/3021
+        self,
+        transport: asyncio.Transport,
+    ):
+        """Method from base class, called when connection is established"""
+        self.transport = transport
+
+    def data_received(self, data: bytes):
+        try:
+            self._data_received(data)
+        except Exception as error:
+            logger.exception("Error receiving data", exc_info=True)
+            self._server._report_server_error(error, JsonRpcInternalError)
+
+    def _data_received(self, data: bytes):
+        """Method from base class, called when server receives the data"""
+        logger.debug("Received %r", data)
+
+        while len(data):
+            # Append the incoming chunk to the message buffer
+            self._message_buf.append(data)
+
+            # Look for the body of the message
+            message = b"".join(self._message_buf)
+            found = JsonRPCProtocol.MESSAGE_PATTERN.fullmatch(message)
+
+            body = found.group("body") if found else b""
+            length = int(found.group("length")) if found else 1
+
+            if len(body) < length:
+                # Message is incomplete; bail until more data arrives
+                return
+
+            # Message is complete;
+            # extract the body and any remaining data,
+            # and reset the buffer for the next message
+            body, data = body[:length], body[length:]
+            self._message_buf = []
+
+            # Parse the body
+            self._procedure_handler(
+                json.loads(
+                    body.decode(self.CHARSET), object_hook=self._deserialize_message
+                )
+            )
+
+    def get_message_type(self, method: str) -> Optional[Type]:
+        """Return the type definition of the message associated with the given method."""
+        return None
+
+    def get_result_type(self, method: str) -> Optional[Type]:
+        """Return the type definition of the result associated with the given method."""
+        return None
+
+    def notify(self, method: str, params=None):
+        """Sends a JSON RPC notification to the client."""
+
+        logger.debug("Sending notification: '%s' %s", method, params)
+
+        notification_type = self.get_message_type(method) or JsonRPCNotification
+        notification = notification_type(
+            method=method, params=params, jsonrpc=JsonRPCProtocol.VERSION
+        )
+
+        self._send_data(notification)
+
+    def send_request(self, method, params=None, callback=None, msg_id=None):
+        """Sends a JSON RPC request to the client.
+
+        Args:
+            method(str): The method name of the message to send
+            params(any): The payload of the message
+
+        Returns:
+            Future that will be resolved once a response has been received
+        """
+
+        if msg_id is None:
+            msg_id = str(uuid.uuid4())
+
+        request_type = self.get_message_type(method) or JsonRPCRequestMessage
+        logger.debug('Sending request with id "%s": %s %s', msg_id, method, params)
+
+        request = request_type(
+            id=msg_id,
+            method=method,
+            params=params,
+            jsonrpc=JsonRPCProtocol.VERSION,
+        )
+
+        future = Future()  # type: ignore[var-annotated]
+        # If callback function is given, call it when result is received
+        if callback:
+
+            def wrapper(future: Future):
+                result = future.result()
+                logger.info("Client response for %s received: %s", params, result)
+                callback(result)
+
+            future.add_done_callback(wrapper)
+
+        self._request_futures[msg_id] = future
+        self._result_types[msg_id] = self.get_result_type(method)
+
+        self._send_data(request)
+
+        return future
+
+    def send_request_async(self, method, params=None, msg_id=None):
+        """Calls `send_request` and wraps `concurrent.futures.Future` with
+        `asyncio.Future` so it can be used with `await` keyword.
+
+        Args:
+            method(str): The method name of the message to send
+            params(any): The payload of the message
+            msg_id(str|int): Optional, message id
+
+        Returns:
+            `asyncio.Future` that can be awaited
+        """
+        return asyncio.wrap_future(
+            self.send_request(method, params=params, msg_id=msg_id)
+        )
+
+    def thread(self):
+        """Decorator that mark function to execute it in a thread."""
+        return self.fm.thread()
+
+
+def lsp_method(method_name: str) -> Callable[[F], F]:
+    def decorator(f: F) -> F:
+        f.method_name = method_name  # type: ignore[attr-defined]
+        return f
+
+    return decorator
+
+
+class LSPMeta(type):
+    """Wraps LSP built-in features (`lsp_` naming convention).
+
+    Built-in features cannot be overridden but user defined features with
+    the same LSP name will be called after them.
+    """
+
+    def __new__(mcs, cls_name, cls_bases, cls):
+        for attr_name, attr_val in cls.items():
+            if callable(attr_val) and hasattr(attr_val, "method_name"):
+                method_name = attr_val.method_name
+                wrapped = call_user_feature(attr_val, method_name)
+                assign_help_attrs(wrapped, method_name, ATTR_FEATURE_TYPE)
+                cls[attr_name] = wrapped
+
+                logger.debug('Added decorator for lsp method: "%s"', attr_name)
+
+        return super().__new__(mcs, cls_name, cls_bases, cls)
+
+
+class LanguageServerProtocol(JsonRPCProtocol, metaclass=LSPMeta):
+    """A class that represents language server protocol.
+
+    It contains implementations for generic LSP features.
+
+    Attributes:
+        workspace(Workspace): In memory workspace
+    """
+
+    def __init__(self, server, converter):
+        super().__init__(server, converter)
+
+        self._workspace: Optional[Workspace] = None
+        self.trace = None
+
+        from jaclang.vendor.pygls.progress import Progress
+
+        self.progress = Progress(self)
+
+        self.server_info = InitializeResultServerInfoType(
+            name=server.name,
+            version=server.version,
+        )
+
+        self._register_builtin_features()
+
+    def _register_builtin_features(self):
+        """Registers generic LSP features from this class."""
+        for name in dir(self):
+            if name in {"workspace"}:
+                continue
+
+            attr = getattr(self, name)
+            if callable(attr) and hasattr(attr, "method_name"):
+                self.fm.add_builtin_feature(attr.method_name, attr)
+
+    @property
+    def workspace(self) -> Workspace:
+        if self._workspace is None:
+            raise RuntimeError(
+                "The workspace is not available - has the server been initialized?"
+            )
+
+        return self._workspace
+
+    @lru_cache()
+    def get_message_type(self, method: str) -> Optional[Type]:
+        """Return LSP type definitions, as provided by `lsprotocol`"""
+        return METHOD_TO_TYPES.get(method, (None,))[0]
+
+    @lru_cache()
+    def get_result_type(self, method: str) -> Optional[Type]:
+        return METHOD_TO_TYPES.get(method, (None, None))[1]
+
+    def apply_edit(
+        self, edit: WorkspaceEdit, label: Optional[str] = None
+    ) -> WorkspaceApplyEditResponse:
+        """Sends apply edit request to the client."""
+        return self.send_request(
+            WORKSPACE_APPLY_EDIT, ApplyWorkspaceEditParams(edit=edit, label=label)
+        )
+
+    def apply_edit_async(
+        self, edit: WorkspaceEdit, label: Optional[str] = None
+    ) -> WorkspaceApplyEditResponse:
+        """Sends apply edit request to the client. Should be called with `await`"""
+        return self.send_request_async(
+            WORKSPACE_APPLY_EDIT, ApplyWorkspaceEditParams(edit=edit, label=label)
+        )
+
+    @lsp_method(EXIT)
+    def lsp_exit(self, *args) -> None:
+        """Stops the server process."""
+        if self.transport is not None:
+            self.transport.close()
+
+        sys.exit(0 if self._shutdown else 1)
+
+    @lsp_method(INITIALIZE)
+    def lsp_initialize(self, params: InitializeParams) -> InitializeResult:
+        """Method that initializes language server.
+        It will compute and return server capabilities based on
+        registered features.
+        """
+        logger.info("Language server initialized %s", params)
+
+        self._server.process_id = params.process_id
+
+        text_document_sync_kind = self._server._text_document_sync_kind
+        notebook_document_sync = self._server._notebook_document_sync
+
+        # Initialize server capabilities
+        self.client_capabilities = params.capabilities
+        self.server_capabilities = ServerCapabilitiesBuilder(
+            self.client_capabilities,
+            set({**self.fm.features, **self.fm.builtin_features}.keys()),
+            self.fm.feature_options,
+            list(self.fm.commands.keys()),
+            text_document_sync_kind,
+            notebook_document_sync,
+        ).build()
+        logger.debug(
+            "Server capabilities: %s",
+            json.dumps(self.server_capabilities, default=self._serialize_message),
+        )
+
+        root_path = params.root_path
+        root_uri = params.root_uri
+        if root_path is not None and root_uri is None:
+            root_uri = from_fs_path(root_path)
+
+        # Initialize the workspace
+        workspace_folders = params.workspace_folders or []
+        self._workspace = Workspace(
+            root_uri,
+            text_document_sync_kind,
+            workspace_folders,
+            self.server_capabilities.position_encoding,
+        )
+
+        self.trace = TraceValues.Off
+
+        return InitializeResult(
+            capabilities=self.server_capabilities,
+            server_info=self.server_info,
+        )
+
+    @lsp_method(INITIALIZED)
+    def lsp_initialized(self, *args) -> None:
+        """Notification received when client and server are connected."""
+        pass
+
+    @lsp_method(SHUTDOWN)
+    def lsp_shutdown(self, *args) -> None:
+        """Request from client which asks server to shutdown."""
+        for future in self._request_futures.values():
+            future.cancel()
+
+        self._shutdown = True
+        return None
+
+    @lsp_method(TEXT_DOCUMENT_DID_CHANGE)
+    def lsp_text_document__did_change(
+        self, params: DidChangeTextDocumentParams
+    ) -> None:
+        """Updates document's content.
+        (Incremental(from server capabilities); not configurable for now)
+        """
+        for change in params.content_changes:
+            self.workspace.update_text_document(params.text_document, change)
+
+    @lsp_method(TEXT_DOCUMENT_DID_CLOSE)
+    def lsp_text_document__did_close(self, params: DidCloseTextDocumentParams) -> None:
+        """Removes document from workspace."""
+        self.workspace.remove_text_document(params.text_document.uri)
+
+    @lsp_method(TEXT_DOCUMENT_DID_OPEN)
+    def lsp_text_document__did_open(self, params: DidOpenTextDocumentParams) -> None:
+        """Puts document to the workspace."""
+        self.workspace.put_text_document(params.text_document)
+
+    @lsp_method(NOTEBOOK_DOCUMENT_DID_OPEN)
+    def lsp_notebook_document__did_open(
+        self, params: DidOpenNotebookDocumentParams
+    ) -> None:
+        """Put a notebook document into the workspace"""
+        self.workspace.put_notebook_document(params)
+
+    @lsp_method(NOTEBOOK_DOCUMENT_DID_CHANGE)
+    def lsp_notebook_document__did_change(
+        self, params: DidChangeNotebookDocumentParams
+    ) -> None:
+        """Update a notebook's contents"""
+        self.workspace.update_notebook_document(params)
+
+    @lsp_method(NOTEBOOK_DOCUMENT_DID_CLOSE)
+    def lsp_notebook_document__did_close(
+        self, params: DidCloseNotebookDocumentParams
+    ) -> None:
+        """Remove a notebook document from the workspace."""
+        self.workspace.remove_notebook_document(params)
+
+    @lsp_method(SET_TRACE)
+    def lsp_set_trace(self, params: SetTraceParams) -> None:
+        """Changes server trace value."""
+        self.trace = params.value
+
+    @lsp_method(WORKSPACE_DID_CHANGE_WORKSPACE_FOLDERS)
+    def lsp_workspace__did_change_workspace_folders(
+        self, params: DidChangeWorkspaceFoldersParams
+    ) -> None:
+        """Adds/Removes folders from the workspace."""
+        logger.info("Workspace folders changed: %s", params)
+
+        added_folders = params.event.added or []
+        removed_folders = params.event.removed or []
+
+        for f_add, f_remove in zip_longest(added_folders, removed_folders):
+            if f_add:
+                self.workspace.add_folder(f_add)
+            if f_remove:
+                self.workspace.remove_folder(f_remove.uri)
+
+    @lsp_method(WORKSPACE_EXECUTE_COMMAND)
+    def lsp_workspace__execute_command(
+        self, params: ExecuteCommandParams, msg_id: str
+    ) -> None:
+        """Executes commands with passed arguments and returns a value."""
+        cmd_handler = self.fm.commands[params.command]
+        self._execute_request(msg_id, cmd_handler, params.arguments)
+
+    @lsp_method(WINDOW_WORK_DONE_PROGRESS_CANCEL)
+    def lsp_work_done_progress_cancel(
+        self, params: WorkDoneProgressCancelParams
+    ) -> None:
+        """Received a progress cancellation from client."""
+        future = self.progress.tokens.get(params.token)
+        if future is None:
+            logger.warning(
+                "Ignoring work done progress cancel for unknown token %s", params.token
+            )
+        else:
+            future.cancel()
+
+    def get_configuration(
+        self,
+        params: WorkspaceConfigurationParams,
+        callback: Optional[ConfigCallbackType] = None,
+    ) -> Future:
+        """Sends configuration request to the client.
+
+        Args:
+            params(WorkspaceConfigurationParams): WorkspaceConfigurationParams from lsp specs
+            callback(callable): Callabe which will be called after
+                                response from the client is received
+        Returns:
+            concurrent.futures.Future object that will be resolved once a
+            response has been received
+        """
+        return self.send_request(WORKSPACE_CONFIGURATION, params, callback)
+
+    def get_configuration_async(
+        self, params: WorkspaceConfigurationParams
+    ) -> asyncio.Future:
+        """Calls `get_configuration` method but designed to use with coroutines
+
+        Args:
+            params(WorkspaceConfigurationParams): WorkspaceConfigurationParams from lsp specs
+        Returns:
+            asyncio.Future that can be awaited
+        """
+        return asyncio.wrap_future(self.get_configuration(params))
+
+    def log_trace(self, message: str, verbose: Optional[str] = None) -> None:
+        """Sends trace notification to the client."""
+        if self.trace == TraceValues.Off:
+            return
+
+        params = LogTraceParams(message=message)
+        if verbose and self.trace == TraceValues.Verbose:
+            params.verbose = verbose
+
+        self.notify(LOG_TRACE, params)
+
+    def _publish_diagnostics_deprecator(
+        self,
+        params_or_uri: Union[str, PublishDiagnosticsParams],
+        diagnostics: Optional[List[Diagnostic]],
+        version: Optional[int],
+        **kwargs,
+    ) -> PublishDiagnosticsParams:
+        if isinstance(params_or_uri, str):
+            message = "DEPRECATION: "
+            "`publish_diagnostics("
+            "self, doc_uri: str, diagnostics: List[Diagnostic], version: Optional[int] = None)`"
+            "will be replaced with `publish_diagnostics(self, params: PublishDiagnosticsParams)`"
+            logging.warning(message)
+
+            params = self._construct_publish_diagnostic_type(
+                params_or_uri, diagnostics, version, **kwargs
+            )
+        else:
+            params = params_or_uri
+        return params
+
+    def _construct_publish_diagnostic_type(
+        self,
+        uri: str,
+        diagnostics: Optional[List[Diagnostic]],
+        version: Optional[int],
+        **kwargs,
+    ) -> PublishDiagnosticsParams:
+        if diagnostics is None:
+            diagnostics = []
+
+        args = {
+            **{"uri": uri, "diagnostics": diagnostics, "version": version},
+            **kwargs,
+        }
+
+        params = PublishDiagnosticsParams(**args)  # type:ignore
+        return params
+
+    def publish_diagnostics(
+        self,
+        params_or_uri: Union[str, PublishDiagnosticsParams],
+        diagnostics: Optional[List[Diagnostic]] = None,
+        version: Optional[int] = None,
+        **kwargs,
+    ):
+        """Sends diagnostic notification to the client.
+
+        .. deprecated:: 1.0.1
+
+           Passing ``(uri, diagnostics, version)`` as arguments is deprecated.
+           Pass an instance of :class:`~lsprotocol.types.PublishDiagnosticParams`
+           instead.
+
+        Parameters
+        ----------
+        params_or_uri
+           The :class:`~lsprotocol.types.PublishDiagnosticParams` to send to the client.
+
+        diagnostics
+           *Deprecated*. The diagnostics to publish
+
+        version
+           *Deprecated*: The version number
+        """
+        params = self._publish_diagnostics_deprecator(
+            params_or_uri, diagnostics, version, **kwargs
+        )
+        self.notify(TEXT_DOCUMENT_PUBLISH_DIAGNOSTICS, params)
+
+    def register_capability(
+        self, params: RegistrationParams, callback: Optional[Callable[[], None]] = None
+    ) -> Future:
+        """Register a new capability on the client.
+
+        Args:
+            params(RegistrationParams): RegistrationParams from lsp specs
+            callback(callable): Callabe which will be called after
+                                response from the client is received
+        Returns:
+            concurrent.futures.Future object that will be resolved once a
+            response has been received
+        """
+        return self.send_request(CLIENT_REGISTER_CAPABILITY, params, callback)
+
+    def register_capability_async(self, params: RegistrationParams) -> asyncio.Future:
+        """Register a new capability on the client.
+
+        Args:
+            params(RegistrationParams): RegistrationParams from lsp specs
+
+        Returns:
+            asyncio.Future object that will be resolved once a
+            response has been received
+        """
+        return asyncio.wrap_future(self.register_capability(params, None))
+
+    def semantic_tokens_refresh(
+        self, callback: Optional[Callable[[], None]] = None
+    ) -> Future:
+        """Requesting a refresh of all semantic tokens.
+
+        Args:
+            callback(callable): Callabe which will be called after
+                                response from the client is received
+
+        Returns:
+            concurrent.futures.Future object that will be resolved once a
+            response has been received
+        """
+        return self.send_request(WORKSPACE_SEMANTIC_TOKENS_REFRESH, callback=callback)
+
+    def semantic_tokens_refresh_async(self) -> asyncio.Future:
+        """Requesting a refresh of all semantic tokens.
+
+        Returns:
+            asyncio.Future object that will be resolved once a
+            response has been received
+        """
+        return asyncio.wrap_future(self.semantic_tokens_refresh(None))
+
+    def show_document(
+        self,
+        params: ShowDocumentParams,
+        callback: Optional[ShowDocumentCallbackType] = None,
+    ) -> Future:
+        """Display a particular document in the user interface.
+
+        Args:
+            params(ShowDocumentParams): ShowDocumentParams from lsp specs
+            callback(callable): Callabe which will be called after
+                                response from the client is received
+
+        Returns:
+            concurrent.futures.Future object that will be resolved once a
+            response has been received
+        """
+        return self.send_request(WINDOW_SHOW_DOCUMENT, params, callback)
+
+    def show_document_async(self, params: ShowDocumentParams) -> asyncio.Future:
+        """Display a particular document in the user interface.
+
+        Args:
+            params(ShowDocumentParams): ShowDocumentParams from lsp specs
+
+        Returns:
+            asyncio.Future object that will be resolved once a
+            response has been received
+        """
+        return asyncio.wrap_future(self.show_document(params, None))
+
+    def show_message(self, message, msg_type=MessageType.Info):
+        """Sends message to the client to display message."""
+        self.notify(
+            WINDOW_SHOW_MESSAGE, ShowMessageParams(type=msg_type, message=message)
+        )
+
+    def show_message_log(self, message, msg_type=MessageType.Log):
+        """Sends message to the client's output channel."""
+        self.notify(
+            WINDOW_LOG_MESSAGE, LogMessageParams(type=msg_type, message=message)
+        )
+
+    def unregister_capability(
+        self,
+        params: UnregistrationParams,
+        callback: Optional[Callable[[], None]] = None,
+    ) -> Future:
+        """Unregister a new capability on the client.
+
+        Args:
+            params(UnregistrationParams): UnregistrationParams from lsp specs
+            callback(callable): Callabe which will be called after
+                                response from the client is received
+        Returns:
+            concurrent.futures.Future object that will be resolved once a
+            response has been received
+        """
+        return self.send_request(CLIENT_UNREGISTER_CAPABILITY, params, callback)
+
+    def unregister_capability_async(
+        self, params: UnregistrationParams
+    ) -> asyncio.Future:
+        """Unregister a new capability on the client.
+
+        Args:
+            params(UnregistrationParams): UnregistrationParams from lsp specs
+            callback(callable): Callabe which will be called after
+                                response from the client is received
+        Returns:
+            asyncio.Future object that will be resolved once a
+            response has been received
+        """
+        return asyncio.wrap_future(self.unregister_capability(params, None))