feat: update server

Author: observerw

src/lsp_client/server/abc.py

@@ -3,13 +3,18 @@
 from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
+from functools import cached_property
 from typing import Self
 
+import anyio
 import asyncer
+from anyio.abc import AnyByteReceiveStream, AnyByteSendStream
+from anyio.streams.buffered import BufferedByteReceiveStream
 from attrs import define, field
 from loguru import logger
 
 from lsp_client.jsonrpc.channel import ResponseTable, response_channel
+from lsp_client.jsonrpc.parse import read_raw_package, write_raw_package
 from lsp_client.jsonrpc.types import (
     RawNotification,
     RawPackage,
@@ -27,51 +32,67 @@ class Server(ABC):
 
     _resp_table: ResponseTable = field(factory=ResponseTable, init=False)
 
+    @property
+    @abstractmethod
+    def send_stream(self) -> AnyByteSendStream:
+        """Stream for sending data to the server."""
+
+    @property
+    @abstractmethod
+    def receive_stream(self) -> AnyByteReceiveStream:
+        """Stream for receiving data from the server."""
+
+    @cached_property
+    def _buffered_receive_stream(self) -> BufferedByteReceiveStream:
+        return BufferedByteReceiveStream(self.receive_stream)
+
     @abstractmethod
     async def check_availability(self) -> None:
         """Check if the server runtime is available."""
 
-    @abstractmethod
+    async def kill(self) -> None:
+        await self.receive_stream.aclose()
+
     async def send(self, package: RawPackage) -> None:
         """Send a package to the runtime."""
+        await write_raw_package(self.send_stream, package)
+        logger.debug("Package sent: {}", package)
 
-    @abstractmethod
     async def receive(self) -> RawPackage | None:
         """Receive a package from the runtime."""
 
-    @abstractmethod
-    async def kill(self) -> None:
-        """Kill the runtime process."""
-
-    async def _dispatch(self, sender: Sender[ServerRequest] | None) -> None:
-        if not sender:
-            logger.warning(
-                "No ServerRequest sender provided, all server requests and notifications will be ignored."
-            )
-
-        async def handle(package: RawPackage) -> None:
-            match package:
-                case {"result": _, "id": id} | {"error": _, "id": id}:
-                    self._resp_table.send(id, package)  # ty: ignore[invalid-argument-type]
-                case {"id": id, "method": _}:
-                    if not sender:
-                        raise RuntimeError(
-                            "Received a server request without a sender provided."
-                        )
-
-                    tx, rx = response_channel.create()
-                    await sender.send((package, tx))  # ty: ignore[invalid-argument-type]
-                    resp = await rx.receive()
-                    await self.send(resp)
-                case {"method": _}:
-                    if not sender:
-                        return
-
-                    await sender.send(package)  # ty: ignore[invalid-argument-type]
-
+        try:
+            package = await read_raw_package(self._buffered_receive_stream)
+            logger.debug("Received package: {}", package)
+            return package
+        except (anyio.EndOfStream, anyio.IncompleteRead, anyio.ClosedResourceError):
+            logger.debug("Stream closed")
+            return None
+
+    async def iter_receive(self) -> AsyncGenerator[RawPackage]:
+        while package := await self.receive():
+            yield package
+
+    async def _handle_package(
+        self, sender: Sender[ServerRequest], package: RawPackage
+    ) -> None:
+        match package:
+            case {"result": _, "id": id} | {"error": _, "id": id} as resp:
+                self._resp_table.send(id, resp)
+            case {"id": id, "method": _} as server_req:
+                tx, rx = response_channel.create()
+                await sender.send((server_req, tx))
+                resp = await rx.receive()
+                await self.send(resp)
+            case {"method": _} as noti:
+                if not sender:
+                    return
+                await sender.send(noti)
+
+    async def _dispatch(self, sender: Sender[ServerRequest]) -> None:
         async with asyncer.create_task_group() as tg:
             while package := await self.receive():
-                tg.soonify(handle)(package)
+                tg.soonify(self._handle_package)(sender, package)
 
     async def request(self, request: RawRequest) -> RawResponsePackage:
         await self.send(request)
@@ -80,12 +101,11 @@ async def request(self, request: RawRequest) -> RawResponsePackage:
     async def notify(self, notification: RawNotification) -> None:
         await self.send(notification)
 
-    @abstractmethod
     @asynccontextmanager
-    def run(
-        self,
-        workspace: Workspace,
-        *,
-        sender: Sender[ServerRequest] | None = None,
+    async def run(
+        self, workspace: Workspace, sender: Sender[ServerRequest]
     ) -> AsyncGenerator[Self]:
         """Run the server."""
+        async with asyncer.create_task_group() as tg:
+            tg.soonify(self._dispatch)(sender)
+            yield self

src/lsp_client/server/container.py

@@ -7,10 +7,10 @@
 from typing import Literal, Self, final, override
 
 import anyio
+from anyio.abc import AnyByteReceiveStream, AnyByteSendStream
 from attrs import Factory, define, field
 from loguru import logger
 
-from lsp_client.jsonrpc.parse import RawPackage
 from lsp_client.server.types import ServerRequest
 from lsp_client.utils.channel import Sender
 from lsp_client.utils.workspace import Workspace
@@ -147,6 +147,34 @@ class ContainerServer(Server):
 
     _local: LocalServer = field(init=False)
 
+    @property
+    @override
+    def send_stream(self) -> AnyByteSendStream:
+        return self._local.send_stream
+
+    @property
+    @override
+    def receive_stream(self) -> AnyByteReceiveStream:
+        return self._local.receive_stream
+
+    @override
+    async def kill(self) -> None:
+        await self._local.kill()
+
+    @override
+    async def check_availability(self) -> None:
+        try:
+            await anyio.run_process(
+                [self.backend, "pull", self.image],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+        except anyio.ProcessError as e:
+            raise RuntimeError(
+                f"Container backend '{self.backend}' is not available or image '{self.image}' "
+                "could not be pulled."
+            ) from e
+
     def format_args(self, workspace: Workspace) -> list[str]:
         args = ["run", "--rm", "-i"]
 
@@ -183,38 +211,9 @@ def format_args(self, workspace: Workspace) -> list[str]:
 
         return args
 
-    @override
-    async def send(self, package: RawPackage) -> None:
-        await self._local.send(package)
-
-    @override
-    async def receive(self) -> RawPackage | None:
-        return await self._local.receive()
-
-    @override
-    async def kill(self) -> None:
-        await self._local.kill()
-
-    @override
-    async def check_availability(self) -> None:
-        try:
-            await anyio.run_process(
-                [self.backend, "pull", self.image],
-                stdout=subprocess.DEVNULL,
-                stderr=subprocess.DEVNULL,
-            )
-        except anyio.ProcessError as e:
-            raise RuntimeError(
-                f"Container backend '{self.backend}' is not available or image '{self.image}' "
-                "could not be pulled."
-            ) from e
-
     @asynccontextmanager
     async def run(
-        self,
-        workspace: Workspace,
-        *,
-        sender: Sender[ServerRequest] | None = None,
+        self, workspace: Workspace, sender: Sender[ServerRequest]
     ) -> AsyncGenerator[Self]:
         args = self.format_args(workspace)
         logger.debug("Running docker runtime with command: {}", args)

src/lsp_client/server/default.py

@@ -1,46 +1,12 @@
 from __future__ import annotations
 
-from collections.abc import AsyncIterator
-from contextlib import suppress
-
 from attrs import frozen
 
-from .abc import Server
 from .container import ContainerServer
-from .error import ServerRuntimeError
 from .local import LocalServer
-from .types import ServerType
 
 
 @frozen
 class DefaultServers:
     local: LocalServer
     container: ContainerServer
-
-    async def iter_candidate(
-        self,
-        *,
-        server: Server | ServerType = "local",
-    ) -> AsyncIterator[Server]:
-        """
-        Server candidates in order of priority:
-        1. User-provided server
-        2. Local server (if available)
-        3. Containerized server
-        4. Local server with auto-install (if enabled)
-        """
-
-        match server:
-            case "container":
-                yield self.container
-            case "local":
-                yield self.local
-            case Server() as server:
-                yield server
-
-        with suppress(ServerRuntimeError):
-            await self.local.check_availability()
-            yield self.local
-
-        yield self.container
-        yield self.local

src/lsp_client/server/local.py

@@ -2,21 +2,16 @@
 
 from collections.abc import AsyncGenerator, Sequence
 from contextlib import asynccontextmanager
-from functools import cached_property
 from pathlib import Path
 from typing import Protocol, Self, final, override
 
 import aioshutil
 import anyio
-import asyncer
-from anyio.abc import AnyByteSendStream, Process
-from anyio.streams.buffered import BufferedByteReceiveStream
+from anyio.abc import AnyByteReceiveStream, AnyByteSendStream, Process
 from attrs import Factory, define, field
 from loguru import logger
 
 from lsp_client.env import disable_auto_installation
-from lsp_client.jsonrpc.parse import read_raw_package, write_raw_package
-from lsp_client.jsonrpc.types import RawPackage
 from lsp_client.server.types import ServerRequest
 from lsp_client.utils.channel import Sender
 from lsp_client.utils.workspace import Workspace
@@ -43,38 +38,25 @@ class LocalServer(Server):
 
     _process: Process = field(init=False, default=None)
 
-    @cached_property
-    def stdin(self) -> AnyByteSendStream:
+    @property
+    @override
+    def send_stream(self) -> AnyByteSendStream:
         stdin = self._process.stdin
         assert stdin, "Process stdin is not available"
         return stdin
 
-    @cached_property
-    def stdout(self) -> BufferedByteReceiveStream:
+    @property
+    @override
+    def receive_stream(self) -> AnyByteReceiveStream:
         stdout = self._process.stdout
         assert stdout, "Process stdout is not available"
-        return BufferedByteReceiveStream(stdout)
+        return stdout
 
-    @cached_property
-    def stderr(self) -> BufferedByteReceiveStream:
+    @property
+    def stderr(self) -> AnyByteReceiveStream:
         stderr = self._process.stderr
         assert stderr, "Process stderr is not available"
-        return BufferedByteReceiveStream(stderr)
-
-    @override
-    async def send(self, package: RawPackage) -> None:
-        await write_raw_package(self.stdin, package)
-        logger.debug("Package sent: {}", package)
-
-    @override
-    async def receive(self) -> RawPackage | None:
-        try:
-            package = await read_raw_package(self.stdout)
-            logger.debug("Received package: {}", package)
-            return package
-        except (anyio.EndOfStream, anyio.IncompleteRead, anyio.ClosedResourceError):
-            logger.debug("Process stdout closed")
-            return None
+        return stderr
 
     @override
     async def kill(self) -> None:
@@ -130,14 +112,10 @@ async def run_process(self, workspace: Workspace) -> AsyncGenerator[None]:
 
     @asynccontextmanager
     async def run(
-        self,
-        workspace: Workspace,
-        *,
-        sender: Sender[ServerRequest] | None = None,
+        self, workspace: Workspace, sender: Sender[ServerRequest]
     ) -> AsyncGenerator[Self]:
         async with (
             self.run_process(workspace),
-            asyncer.create_task_group() as tg,
+            super().run(workspace, sender=sender),
         ):
-            tg.soonify(self._dispatch)(sender)
             yield self