Repository: openlawlibrary/pygls
Pull request #501 (2024-10-07 20:09:49)

Add TCP support to pygls' LanguageClient
## Description (e.g. "Related to ...", etc.)

This PR adds a `start_tcp` method to pygls' `LanguageClient`. This method uses the high-level asyncio APIs I would like to transition the server side over to using (as mentioned in #334).

However, before touching the server-side code I thought it would be best to add tests around the current implementation. So this PR also extends pygls' test suite to add a `--lsp-transport` command line argument.

- `--lsp-transport stdio` (the default), runs the end-to-end tests over stdin/stdout.
- `--lsp-transport tcp`, runs the end-to-end tests over a TCP connection

The `poe test` task used in CI has been updated to run both style of end-to-end tests.

Finaly, this PR also includes some smaller client related fixes
- When the server process exits (in stdio mode), any pending requests are automatically cancelled
- In stdio mode, the client waits for the server process to exit rather than agressively terminating it (fixing coverage reporting)

## Code review checklist (for code reviewer to complete)

- [ ] Pull request represents a single change (i.e. not fixing disparate/unrelated things in a single PR)
- [ ] Title summarizes what is changing
- [ ] Commit messages are meaningful (see [this][commit messages] for details)
- [ ] Tests have been included and/or updated, as appropriate
- [ ] Docstrings have been included and/or updated, as appropriate
- [ ] Standalone docs have been updated accordingly

## Automated linters

You can run the lints that are run on CI locally with:
```sh
poetry install --all-extras --with dev
poetry run poe lint
```

[commit messages]: https://conventionalcommits.org/


alcarney: Running the end-to-end tests over TCP has highlighted an interesting bug (#502). Since it isn't really related to the client's new TCP capability I've skipped the test for now and we can figure out a fix for it separately - who knows maybe future server side refactorings will magically fix it? :sweat_smile: :pray: 
alcarney: Thanks!

Patch:
diff --git a/docs/source/examples/json-server.rst b/docs/source/examples/json-server.rst
index 1c2f483..6f236bc 100644
--- a/docs/source/examples/json-server.rst
+++ b/docs/source/examples/json-server.rst
@@ -2,12 +2,4 @@ JSON Server
 ===========
 
 .. example-server:: json_server.py
-   :start-at: import argparse
-
-
-
-
-
-
-
-
+   :start-at: import asyncio
diff --git a/examples/servers/code_actions.py b/examples/servers/code_actions.py
index 26fe1ad..bb1bdbc 100644
--- a/examples/servers/code_actions.py
+++ b/examples/servers/code_actions.py
@@ -26,6 +26,7 @@ action which, when invoked will fill in the answer.
 """
 
 import re
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 from lsprotocol import types
 
@@ -74,4 +75,4 @@ def code_actions(params: types.CodeActionParams):
 
 
 if __name__ == "__main__":
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/code_lens.py b/examples/servers/code_lens.py
index dcbc7aa..f7b41ba 100644
--- a/examples/servers/code_lens.py
+++ b/examples/servers/code_lens.py
@@ -34,6 +34,7 @@ import re
 
 from lsprotocol import types
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 
 ADDITION = re.compile(r"^\s*(\d+)\s*\+\s*(\d+)\s*=(?=\s*$)")
@@ -140,4 +141,4 @@ def evaluate_sum(ls: LanguageServer, args):
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, format="%(message)s")
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/colors.py b/examples/servers/colors.py
index 28884c9..b1040ce 100644
--- a/examples/servers/colors.py
+++ b/examples/servers/colors.py
@@ -49,6 +49,7 @@ import re
 
 from lsprotocol import types
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 
 COLOR = re.compile(r"""\#([a-fA-F0-9]{6}|[a-fA-F0-9]{3})(?!\w)""")
@@ -112,4 +113,4 @@ def color_presentation(params: types.ColorPresentationParams):
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, format="%(message)s")
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/formatting.py b/examples/servers/formatting.py
index 5c736d2..1e65604 100644
--- a/examples/servers/formatting.py
+++ b/examples/servers/formatting.py
@@ -42,6 +42,7 @@ from typing import Optional
 import attrs
 from lsprotocol import types
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 from pygls.workspace import TextDocument
 
@@ -194,5 +195,4 @@ def skip_line(line: int, range_: Optional[types.Range]) -> bool:
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, format="%(message)s")
-
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/goto.py b/examples/servers/goto.py
index f2e081b..cdc7694 100644
--- a/examples/servers/goto.py
+++ b/examples/servers/goto.py
@@ -38,6 +38,7 @@ import re
 
 from lsprotocol import types
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 from pygls.workspace import TextDocument
 
@@ -210,4 +211,4 @@ def find_references(ls: GotoLanguageServer, params: types.ReferenceParams):
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, format="%(message)s")
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/hover.py b/examples/servers/hover.py
index 90aa3e5..68bf63e 100644
--- a/examples/servers/hover.py
+++ b/examples/servers/hover.py
@@ -29,6 +29,7 @@ from datetime import datetime
 
 from lsprotocol import types
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 
 DATE_FORMATS = [
@@ -84,5 +85,4 @@ def hover(ls: LanguageServer, params: types.HoverParams):
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, format="%(message)s")
-
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/inlay_hints.py b/examples/servers/inlay_hints.py
index 4e96109..21cb56e 100644
--- a/examples/servers/inlay_hints.py
+++ b/examples/servers/inlay_hints.py
@@ -33,6 +33,7 @@ import re
 from typing import Optional
 
 from lsprotocol import types
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 
 NUMBER = re.compile(r"\d+")
@@ -91,4 +92,4 @@ def inlay_hint_resolve(hint: types.InlayHint):
 
 
 if __name__ == "__main__":
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/json_server.py b/examples/servers/json_server.py
index be0641e..389bb6d 100644
--- a/examples/servers/json_server.py
+++ b/examples/servers/json_server.py
@@ -30,7 +30,6 @@ Eventually this example will be broken up in smaller, more focused examples and
 guides.
 """
 
-import argparse
 import asyncio
 import uuid
 from functools import partial
@@ -38,6 +37,7 @@ from typing import Optional
 
 from lsprotocol import types as lsp
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 
 
@@ -232,27 +232,5 @@ def show_configuration_thread(ls: JsonLanguageServer, *args):
     handle_config(ls, config)
 
 
-def add_arguments(parser):
-    parser.description = "simple json server example"
-
-    parser.add_argument("--tcp", action="store_true", help="Use TCP server")
-    parser.add_argument("--ws", action="store_true", help="Use WebSocket server")
-    parser.add_argument("--host", default="127.0.0.1", help="Bind to this address")
-    parser.add_argument("--port", type=int, default=2087, help="Bind to this port")
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    add_arguments(parser)
-    args = parser.parse_args()
-
-    if args.tcp:
-        json_server.start_tcp(args.host, args.port)
-    elif args.ws:
-        json_server.start_ws(args.host, args.port)
-    else:
-        json_server.start_io()
-
-
 if __name__ == "__main__":
-    main()
+    start_server(json_server)
diff --git a/examples/servers/links.py b/examples/servers/links.py
index 99bacdd..dd81c3c 100644
--- a/examples/servers/links.py
+++ b/examples/servers/links.py
@@ -33,6 +33,7 @@ import re
 
 from lsprotocol import types
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 
 LINK = re.compile(r"<(\w+):([^>]+)>")
@@ -91,4 +92,4 @@ def document_link_resolve(link: types.DocumentLink):
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, format="%(message)s")
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/publish_diagnostics.py b/examples/servers/publish_diagnostics.py
index 3eb8ef6..7b14d87 100644
--- a/examples/servers/publish_diagnostics.py
+++ b/examples/servers/publish_diagnostics.py
@@ -30,6 +30,7 @@ import re
 
 from lsprotocol import types
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 from pygls.workspace import TextDocument
 
@@ -117,4 +118,4 @@ def did_change(ls: PublishDiagnosticServer, params: types.DidOpenTextDocumentPar
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, format="%(message)s")
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/pull_diagnostics.py b/examples/servers/pull_diagnostics.py
index b33e77f..e51aa0a 100644
--- a/examples/servers/pull_diagnostics.py
+++ b/examples/servers/pull_diagnostics.py
@@ -32,6 +32,7 @@ import re
 
 from lsprotocol import types
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 from pygls.workspace import TextDocument
 
@@ -160,4 +161,4 @@ def workspace_diagnostic(
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, format="%(message)s")
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/rename.py b/examples/servers/rename.py
index e0c1fc2..8343713 100644
--- a/examples/servers/rename.py
+++ b/examples/servers/rename.py
@@ -36,6 +36,7 @@ from typing import List
 
 from lsprotocol import types
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 from pygls.workspace import TextDocument
 
@@ -155,4 +156,4 @@ def prepare_rename(ls: RenameLanguageServer, params: types.PrepareRenameParams):
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.DEBUG, format="%(message)s")
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/semantic_tokens.py b/examples/servers/semantic_tokens.py
index 7c44618..194eea6 100644
--- a/examples/servers/semantic_tokens.py
+++ b/examples/servers/semantic_tokens.py
@@ -42,6 +42,7 @@ from typing import Optional
 import attrs
 from lsprotocol import types
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 from pygls.workspace import TextDocument
 
@@ -299,5 +300,4 @@ def semantic_tokens_full(ls: SemanticTokensServer, params: types.SemanticTokensP
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, format="%(message)s")
-
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/symbols.py b/examples/servers/symbols.py
index 4d0a7fb..1e4c4a2 100644
--- a/examples/servers/symbols.py
+++ b/examples/servers/symbols.py
@@ -39,6 +39,7 @@ import re
 
 from lsprotocol import types
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 from pygls.workspace import TextDocument
 
@@ -245,4 +246,4 @@ def workspace_symbol(ls: SymbolsLanguageServer, params: types.WorkspaceSymbolPar
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, format="%(message)s")
-    server.start_io()
+    start_server(server)
diff --git a/examples/servers/threaded_handlers.py b/examples/servers/threaded_handlers.py
index 8d385fe..fba9fcd 100644
--- a/examples/servers/threaded_handlers.py
+++ b/examples/servers/threaded_handlers.py
@@ -25,6 +25,7 @@ import threading
 
 from lsprotocol import types
 
+from pygls.cli import start_server
 from pygls.lsp.server import LanguageServer
 
 server = LanguageServer("threaded-server", "v1")
@@ -92,4 +93,4 @@ def count_down_error(ls: LanguageServer, *args):
 
 
 if __name__ == "__main__":
-    server.start_io()
+    start_server(server)
diff --git a/poetry.lock b/poetry.lock
index 50c9213..3aa1eba 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -916,6 +916,24 @@ pytest = ">=7.0.0,<9"
 docs = ["sphinx (>=5.3)", "sphinx-rtd-theme (>=1.0)"]
 testing = ["coverage (>=6.2)", "hypothesis (>=5.7.1)"]
 
+[[package]]
+name = "pytest-cov"
+version = "5.0.0"
+description = "Pytest plugin for measuring coverage."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "pytest-cov-5.0.0.tar.gz", hash = "sha256:5837b58e9f6ebd335b0f8060eecce69b662415b16dc503883a02f45dfeb14857"},
+    {file = "pytest_cov-5.0.0-py3-none-any.whl", hash = "sha256:4f0764a1219df53214206bf1feea4633c3b558a2925c8b59f144f682861ce652"},
+]
+
+[package.dependencies]
+coverage = {version = ">=5.2.1", extras = ["toml"]}
+pytest = ">=4.6"
+
+[package.extras]
+testing = ["fields", "hunter", "process-tests", "pytest-xdist", "virtualenv"]
+
 [[package]]
 name = "pyyaml"
 version = "6.0.2"
@@ -1468,4 +1486,4 @@ ws = ["websockets"]
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9"
-content-hash = "1dde66e7780b345c31879a54f573c8c2ea813177df09f5e3610395e9c91c0976"
+content-hash = "d914d6b11699574685a04f4bd1201031b7bc1485497c5e5656454120856692c9"
diff --git a/pygls/cli.py b/pygls/cli.py
new file mode 100644
index 0000000..bfcb4d8
--- /dev/null
+++ b/pygls/cli.py
@@ -0,0 +1,45 @@
+############################################################################
+# 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.                                           #
+############################################################################
+"""A simple cli wrapper for pygls servers."""
+from __future__ import annotations
+
+import argparse
+import typing
+
+if typing.TYPE_CHECKING:
+    from pygls.server import JsonRPCServer
+
+
+def start_server(server: JsonRPCServer, args: list[str] | None = None):
+    """A helper function that implements a simple cli wrapper for a pygls server
+    allowing the user to select between the supported transports."""
+
+    name = type(server).__name__
+    parser = argparse.ArgumentParser(description=f"start a {name} instance")
+    parser.add_argument("--tcp", action="store_true", help="start a TCP server")
+    parser.add_argument("--ws", action="store_true", help="start a WebSocket server")
+    parser.add_argument("--host", default="127.0.0.1", help="bind to this address")
+    parser.add_argument("--port", type=int, default=8888, help="bind to this port")
+
+    arguments = parser.parse_args(args)
+
+    if arguments.tcp:
+        server.start_tcp(arguments.host, arguments.port)
+    elif arguments.ws:
+        server.start_ws(arguments.host, arguments.port)
+    else:
+        server.start_io()
diff --git a/pygls/client.py b/pygls/client.py
index 577f05e..b99be05 100644
--- a/pygls/client.py
+++ b/pygls/client.py
@@ -14,61 +14,29 @@
 # See the License for the specific language governing permissions and      #
 # limitations under the License.                                           #
 ############################################################################
+from __future__ import annotations
+
 import asyncio
+import json
 import logging
 import re
+import typing
 from threading import Event
-from typing import Any
-from typing import Callable
-from typing import List
-from typing import Optional
-from typing import Type
-from typing import Union
-
-from cattrs import Converter
 
-from pygls.exceptions import PyglsError, JsonRpcException
+from pygls.exceptions import PyglsError, JsonRpcException, JsonRpcInternalError
 from pygls.protocol import JsonRPCProtocol, default_converter
 
+if typing.TYPE_CHECKING:
+    from typing import Any
+    from typing import Callable
+    from typing import List
+    from typing import Optional
+    from typing import Type
 
-logger = logging.getLogger(__name__)
+    from cattrs import Converter
 
 
-async def aio_readline(stop_event, reader, message_handler):
-    CONTENT_LENGTH_PATTERN = re.compile(rb"^Content-Length: (\d+)\r\n$")
-
-    # Initialize message buffer
-    message = []
-    content_length = 0
-
-    while not stop_event.is_set():
-        # Read a header line
-        header = await reader.readline()
-        if not header:
-            break
-        message.append(header)
-
-        # Extract content length if possible
-        if not content_length:
-            match = CONTENT_LENGTH_PATTERN.fullmatch(header)
-            if match:
-                content_length = int(match.group(1))
-                logger.debug("Content length: %s", content_length)
-
-        # Check if all headers have been read (as indicated by an empty line \r\n)
-        if content_length and not header.strip():
-            # Read body
-            body = await reader.readexactly(content_length)
-            if not body:
-                break
-            message.append(body)
-
-            # Pass message to protocol
-            message_handler(b"".join(message))
-
-            # Reset the buffer
-            message = []
-            content_length = 0
+logger = logging.getLogger(__name__)
 
 
 class JsonRPCClient:
@@ -79,7 +47,7 @@ class JsonRPCClient:
         protocol_cls: Type[JsonRPCProtocol] = JsonRPCProtocol,
         converter_factory: Callable[[], Converter] = default_converter,
     ):
-        # Strictly speaking `JsonRPCProtocol` wants a `LanguageServer`, not a
+        # Strictly speaking `JsonRPCProtocol` wants a `JsonRPCServer`, not a
         # `JsonRPCClient`. However there similar enough for our purposes, which is
         # that this client will mostly be used in testing contexts.
         self.protocol = protocol_cls(self, converter_factory())  # type: ignore
@@ -128,39 +96,103 @@ class JsonRPCClient:
             **kwargs,
         )
 
+        # Keep mypy happy
+        if server.stdout is None:
+            raise RuntimeError("Server process is missing a stdout stream")
+
         self.protocol.connection_made(server.stdin)  # type: ignore
-        connection = asyncio.create_task(
-            aio_readline(self._stop_event, server.stdout, self.protocol.data_received)
-        )
+        connection = asyncio.create_task(self.run_async(server.stdout))
         notify_exit = asyncio.create_task(self._server_exit())
 
         self._server = server
         self._async_tasks.extend([connection, notify_exit])
 
+    async def start_tcp(self, host: str, port: int):
+        """Start communicating with a server over TCP."""
+        reader, writer = await asyncio.open_connection(host, port)
+
+        self.protocol.connection_made(writer)  # type: ignore
+        connection = asyncio.create_task(self.run_async(reader))
+
+        self._async_tasks.extend([connection])
+
+    async def run_async(self, reader: asyncio.StreamReader):
+        """Run the main message processing loop, asynchronously"""
+
+        CONTENT_LENGTH_PATTERN = re.compile(rb"^Content-Length: (\d+)\r\n$")
+        content_length = 0
+
+        while not self._stop_event.is_set():
+            # Read a header line
+            header = await reader.readline()
+            if not header:
+                break
+
+            # Extract content length if possible
+            if not content_length:
+                match = CONTENT_LENGTH_PATTERN.fullmatch(header)
+                if match:
+                    content_length = int(match.group(1))
+                    logger.debug("Content length: %s", content_length)
+
+            # Check if all headers have been read (as indicated by an empty line \r\n)
+            if content_length and not header.strip():
+                # Read body
+                body = await reader.readexactly(content_length)
+                if not body:
+                    break
+
+                try:
+                    message = json.loads(
+                        body, object_hook=self.protocol.structure_message
+                    )
+                    self.protocol.handle_message(message)
+                except Exception as exc:
+                    logger.exception("Unable to handle message")
+                    self._report_server_error(exc, JsonRpcInternalError)
+
+                # Reset
+                content_length = 0
+
     async def _server_exit(self):
-        if self._server is not None:
-            await self._server.wait()
-            logger.debug(
-                "Server process %s exited with return code: %s",
-                self._server.pid,
-                self._server.returncode,
-            )
+        """Cleanup handler that runs when the server process managed by the client exits"""
+        if self._server is None:
+            return
+
+        await self._server.wait()
+
+        pid = self._server.pid
+        returncode = self._server.returncode
+
+        reason = f"Server process {pid} exited with return code: {returncode}"
+        logger.debug(reason)
+
+        # Cancel any pending requests
+        for id_, fut in self.protocol._request_futures.items():
+            if not fut.done():
+                fut.set_exception(RuntimeError(reason))
+                logger.debug("Cancelled pending request '%s': %s", id_, reason)
+
+        try:
             await self.server_exit(self._server)
-            self._stop_event.set()
+        except Exception:
+            logger.exception("Error in server_exit handler")
+
+        self._stop_event.set()
 
     async def server_exit(self, server: asyncio.subprocess.Process):
         """Called when the server process exits."""
 
     def _report_server_error(
-        self, error: Exception, source: Union[PyglsError, JsonRpcException]
+        self, error: Exception, source: type[PyglsError] | type[JsonRpcException]
     ):
         try:
             self.report_server_error(error, source)
         except Exception:
-            logger.error("Unable to report error", exc_info=True)
+            logger.exception("Unable to report error")
 
     def report_server_error(
-        self, error: Exception, source: Union[PyglsError, JsonRpcException]
+        self, error: Exception, source: type[PyglsError] | type[JsonRpcException]
     ):
         """Called when the server does something unexpected e.g. respond with malformed
         JSON."""
@@ -169,8 +201,7 @@ class JsonRPCClient:
         self._stop_event.set()
 
         if self._server is not None and self._server.returncode is None:
-            logger.debug("Terminating server process: %s", self._server.pid)
-            self._server.terminate()
+            await self._server.wait()
 
         if len(self._async_tasks) > 0:
             await asyncio.gather(*self._async_tasks)
diff --git a/pygls/protocol/json_rpc.py b/pygls/protocol/json_rpc.py
index b568eb6..952a7c6 100644
--- a/pygls/protocol/json_rpc.py
+++ b/pygls/protocol/json_rpc.py
@@ -312,7 +312,7 @@ class JsonRPCProtocol(asyncio.Protocol):
 
         return data.__dict__
 
-    def _deserialize_message(self, data):
+    def structure_message(self, data):
         """Function used to deserialize data recevied from the client."""
 
         if "jsonrpc" not in data:
@@ -346,7 +346,7 @@ class JsonRPCProtocol(asyncio.Protocol):
             logger.error("Unable to deserialize message\n%s", traceback.format_exc())
             raise JsonRpcInternalError() from exc
 
-    def _procedure_handler(self, message):
+    def handle_message(self, message):
         """Delegates message to handlers depending on message type."""
 
         if message.jsonrpc != JsonRPCProtocol.VERSION:
@@ -470,9 +470,9 @@ class JsonRPCProtocol(asyncio.Protocol):
             self._message_buf = []
 
             # Parse the body
-            self._procedure_handler(
+            self.handle_message(
                 json.loads(
-                    body.decode(self.CHARSET), object_hook=self._deserialize_message
+                    body.decode(self.CHARSET), object_hook=self.structure_message
                 )
             )
 
diff --git a/pygls/server.py b/pygls/server.py
index 22fd2a7..c51f909 100644
--- a/pygls/server.py
+++ b/pygls/server.py
@@ -292,8 +292,8 @@ class JsonRPCServer:
             """Handle new connection wrapped in the WebSocket."""
             self.protocol.transport = WebSocketTransportAdapter(websocket, self.loop)
             async for message in websocket:
-                self.protocol._procedure_handler(
-                    json.loads(message, object_hook=self.protocol._deserialize_message)
+                self.protocol.handle_message(
+                    json.loads(message, object_hook=self.protocol.structure_message)
                 )
 
         start_server = serve(connection_made, host, port, loop=self.loop)
diff --git a/pyproject.toml b/pyproject.toml
index 4cb365a..a8b6ea0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,6 +38,7 @@ black = "^24.4.1"
 coverage = { version = ">=7.3.2", extras = ["toml"] }
 pytest = ">=7.4.3"
 pytest-asyncio = ">=0.21.0"
+pytest-cov = ">=5"
 
 [tool.poetry.group.docs.dependencies]
 # TODO `sphinx>=7.26` needs python 3.9
@@ -64,13 +65,9 @@ black_check = "black --check ."
 poetry_lock_check = "poetry check"
 
 [tool.poe.tasks.test]
-env = { COVERAGE_PROCESS_START = "pyproject.toml" }
 sequence = [
-     { cmd = "python -c 'import pathlib,sys; pathlib.Path(f\"{sys.path[-1]}/cov.pth\").write_text(\"import coverage; coverage.process_startup()\");'"},
-     { cmd = "coverage erase" },
-     { cmd = "coverage run -m pytest" },
-     { cmd = "coverage combine" },
-     { cmd = "coverage report" },
+     { cmd = "pytest --cov" },
+     { cmd = "pytest tests/e2e --lsp-transport tcp" },
 ]
 ignore_fail = "return_non_zero"
 


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 2e89d2e..fb896f2 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -16,6 +16,8 @@
 # See the License for the specific language governing permissions and      #
 # limitations under the License.                                           #
 ############################################################################
+from __future__ import annotations
+
 import asyncio
 import pathlib
 import sys
@@ -126,6 +128,15 @@ def pytest_addoption(parser):
         help="Choose the runtime in which to run servers under test.",
     )
 
+    group.addoption(
+        "--lsp-transport",
+        dest="lsp_transport",
+        action="store",
+        default="stdio",
+        choices=("stdio", "tcp"),
+        help="Choose the transport to use with servers under test.",
+    )
+
 
 @pytest.fixture(scope="session")
 def runtime(request):
@@ -134,6 +145,13 @@ def runtime(request):
     return request.config.getoption("lsp_runtime")
 
 
+@pytest.fixture(scope="session")
+def transport(request):
+    """This fixture is the source of truth for the transport we should run the
+    end-to-end tests with."""
+    return request.config.getoption("lsp_transport")
+
+
 @pytest.fixture(scope="session")
 def path_for():
     """Returns the path corresponding to a file in the example workspace"""
@@ -171,14 +189,31 @@ def server_dir():
     return path.resolve()
 
 
-def get_client_for_cpython_server(uri_fixture):
+def get_client_for_cpython_server(transport, uri_fixture):
     """Return a client configured to communicate with a server running under cpython."""
 
     async def fn(
         server_name: str, capabilities: Optional[types.ClientCapabilities] = None
     ):
         client = LanguageClient("pygls-test-suite", "v1")
-        await client.start_io(sys.executable, str(SERVER_DIR / server_name))
+
+        server_cmd = [sys.executable, str(SERVER_DIR / server_name)]
+        server: asyncio.subprocess.Process | None = None
+
+        if transport == "stdio":
+            await client.start_io(*server_cmd)
+
+        elif transport == "tcp":
+            # TODO: Make host/port configurable?
+            host, port = "localhost", 8888
+            server_cmd.extend(["--tcp", "--host", host, "--port", f"{port}"])
+
+            server = await asyncio.create_subprocess_exec(*server_cmd)
+            await asyncio.sleep(1)
+            await client.start_tcp(host, port)
+
+        else:
+            raise NotImplementedError(f"Unsupported transport: {transport!r}")
 
         response = await client.initialize_async(
             types.InitializeParams(
@@ -193,19 +228,19 @@ def get_client_for_cpython_server(uri_fixture):
         client.exit(None)
 
         await client.stop()
+        if server is not None and server.returncode is None:
+            server.terminate()
 
     return fn
 
 
 @pytest.fixture(scope="session")
-def get_client_for(runtime, uri_for):
+def get_client_for(runtime, transport, uri_for):
     """Return a client configured to communicate with the specified server.
 
-    Takes into account the current runtime.
-
-    It's the consuming fixture's responsibility to stop the client.
+    Takes into account the current runtime and transport.
     """
     if runtime not in {"cpython"}:
         raise NotImplementedError(f"get_client_for: {runtime=}")
 
-    return get_client_for_cpython_server(uri_for)
+    return get_client_for_cpython_server(transport, uri_for)
diff --git a/tests/e2e/test_threaded_handlers.py b/tests/e2e/test_threaded_handlers.py
index 54f068e..74077a8 100644
--- a/tests/e2e/test_threaded_handlers.py
+++ b/tests/e2e/test_threaded_handlers.py
@@ -24,6 +24,7 @@ from functools import partial
 import pytest
 import pytest_asyncio
 from lsprotocol import types
+from pygls import IS_WIN
 from pygls.exceptions import JsonRpcInternalError
 
 if typing.TYPE_CHECKING:
@@ -123,9 +124,15 @@ async def test_countdown_blocking(
 
 @pytest.mark.asyncio(scope="function")
 async def test_countdown_threaded(
-    threaded_handlers: Tuple[BaseLanguageClient, types.InitializeResult], uri_for
+    threaded_handlers: Tuple[BaseLanguageClient, types.InitializeResult],
+    uri_for,
+    transport,
 ):
     """Ensure that the countdown threaded command is working as expected."""
+
+    if IS_WIN and transport == "tcp":
+        pytest.skip("see https://github.com/openlawlibrary/pygls/issues/502")
+
     client, initialize_result = threaded_handlers
 
     completion_options = initialize_result.capabilities.completion_provider
diff --git a/tests/servers/large_response.py b/tests/servers/large_response.py
index 04cba3a..61e1990 100644
--- a/tests/servers/large_response.py
+++ b/tests/servers/large_response.py
@@ -1,37 +1,22 @@
 """This server returns a particuarly large response."""
 
-import asyncio
-import threading
 import sys
-from concurrent.futures import ThreadPoolExecutor
 
-from pygls.server import aio_readline
+from pygls.server import JsonRPCServer
+from pygls.protocol import JsonRPCProtocol, default_converter
 
+server = JsonRPCServer(JsonRPCProtocol, default_converter)
 
-def handler(data):
-    payload = dict(
-        jsonrpc="2.0",
-        id=1,
-        result=dict(
-            numbers=list(range(100_000)),
-        ),
-    )
-    content = str(payload).replace("'", '"')
-    message = f"Content-Length: {len(content)}\r\n\r\n{content}".encode("utf8")
 
-    sys.stdout.buffer.write(message)
-    sys.stdout.flush()
+@server.feature("get/numbers")
+def get_numbers(*args):
+    return dict(numbers=list(range(100_000)))
 
 
-async def main():
-    await aio_readline(
-        asyncio.get_running_loop(),
-        ThreadPoolExecutor(),
-        threading.Event(),
-        sys.stdin.buffer,
-        handler,
-    )
+@server.feature("exit")
+def exit(*args):
+    sys.exit(0)
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    server.start_io()
diff --git a/tests/test_client.py b/tests/test_client.py
index cacd3e0..8fb1327 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -99,4 +99,5 @@ async def test_client_large_responses():
     result = await client.protocol.send_request_async("get/numbers", {}, msg_id=1)
     assert len(result.numbers) == 100_000
 
-    await client.stop()
+    client.protocol.notify("exit", {})
+    await asyncio.wait_for(client.stop(), timeout=5.0)
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index bc73e37..daabaec 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -115,7 +115,7 @@ def test_deserialize_notification_message_valid_params(protocol):
     }}
     """
 
-    result = json.loads(params, object_hook=protocol._deserialize_message)
+    result = json.loads(params, object_hook=protocol.structure_message)
 
     assert isinstance(
         result, ExampleNotification
@@ -144,7 +144,7 @@ def test_deserialize_notification_message_unknown_type(protocol):
     }
     """
 
-    result = json.loads(params, object_hook=protocol._deserialize_message)
+    result = json.loads(params, object_hook=protocol.structure_message)
 
     assert isinstance(result, JsonRPCNotification)
     assert result.jsonrpc == "2.0"
@@ -169,7 +169,7 @@ def test_deserialize_notification_message_bad_params_should_raise_error(protocol
     """
 
     with pytest.raises(JsonRpcInvalidParams):
-        json.loads(params, object_hook=protocol._deserialize_message)
+        json.loads(params, object_hook=protocol.structure_message)
 
 
 def test_deserialize_response_message_custom_converter():
@@ -200,7 +200,7 @@ def test_deserialize_response_message_custom_converter():
 
     protocol = JsonRPCProtocol(None, custom_converter())
     protocol._result_types["id"] = egasseM
-    result = json.loads(params, object_hook=protocol._deserialize_message)
+    result = json.loads(params, object_hook=protocol.structure_message)
 
     assert isinstance(result, egasseM)
     assert result.cprnosj == "2.0"
@@ -296,7 +296,7 @@ def test_deserialize_response_message(protocol):
     }
     """
     protocol._result_types["id"] = IntResult
-    result = json.loads(params, object_hook=protocol._deserialize_message)
+    result = json.loads(params, object_hook=protocol.structure_message)
 
     assert isinstance(result, IntResult)
     assert result.jsonrpc == "2.0"
@@ -318,7 +318,7 @@ def test_deserialize_response_message_unknown_type(protocol):
     }
     """
     protocol._result_types["id"] = JsonRPCResponseMessage
-    result = json.loads(params, object_hook=protocol._deserialize_message)
+    result = json.loads(params, object_hook=protocol.structure_message)
 
     assert isinstance(result, JsonRPCResponseMessage)
     assert result.jsonrpc == "2.0"
@@ -342,7 +342,7 @@ def test_deserialize_request_message_with_registered_type(protocol):
         }}
     }}
     """
-    result = json.loads(params, object_hook=protocol._deserialize_message)
+    result = json.loads(params, object_hook=protocol.structure_message)
 
     assert isinstance(result, ExampleRequest)
     assert result.jsonrpc == "2.0"
@@ -370,7 +370,7 @@ def test_deserialize_request_message_without_registered_type(protocol):
         }
     }
     """
-    result = json.loads(params, object_hook=protocol._deserialize_message)
+    result = json.loads(params, object_hook=protocol.structure_message)
 
     assert isinstance(result, JsonRPCRequestMessage)
     assert result.jsonrpc == "2.0"