hermes-agent/agent/lsp/protocol.py

"""Minimal LSP JSON-RPC 2.0 framer over async streams.

Wire format: ``Content-Length: <bytes>\\r\\n\\r\\n<utf-8 JSON body>`` where the
body is a JSON-RPC 2.0 request, response, or notification.  Just the framer
plus envelope helpers, so :class:`agent.lsp.client.LSPClient` can focus on
protocol semantics.
"""
from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, Optional, Tuple

logger = logging.getLogger("agent.lsp.protocol")

# LSP error codes we care about (spec 3.17 #errorCodes).
ERROR_CONTENT_MODIFIED = -32801
ERROR_REQUEST_CANCELLED = -32800
ERROR_METHOD_NOT_FOUND = -32601

_MAX_HEADER_BYTES = 8192  # a well-behaved server fits in well under 200 bytes
_MAX_BODY_BYTES = 64 * 1024 * 1024


class LSPProtocolError(Exception):
    """The framing or envelope itself is broken (vs. :class:`LSPRequestError`, a conformant error response)."""


class LSPRequestError(Exception):
    """An LSP request returned a JSON-RPC error response; carries ``code``, ``message``, ``data``."""

    def __init__(self, code: int, message: str, data: Any = None) -> None:
        super().__init__(f"LSP error {code}: {message}")
        self.code = code
        self.message = message
        self.data = data


def encode_message(obj: dict) -> bytes:
    """Encode an envelope as compact UTF-8 JSON with an exact Content-Length header."""
    body = json.dumps(obj, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    header = f"Content-Length: {len(body)}\r\n\r\n".encode("ascii")
    return header + body


async def read_message(reader: asyncio.StreamReader) -> Optional[dict]:
    """Read one framed message.

    Returns ``None`` on clean EOF between messages (typical shutdown);
    raises :class:`LSPProtocolError` on malformed framing.
    """
    headers: dict = {}
    header_bytes = 0
    while True:
        try:
            line = await reader.readuntil(b"\r\n")
        except asyncio.IncompleteReadError as e:
            # EOF before any header started is a clean close; mid-block is bad framing.
            if not e.partial and not headers:
                return None
            raise LSPProtocolError(
                f"unexpected EOF while reading LSP headers (partial={e.partial!r})"
            ) from e
        # Cap against a server streaming headers without ever emitting CRLF-CRLF.
        header_bytes += len(line)
        if header_bytes > _MAX_HEADER_BYTES:
            raise LSPProtocolError(
                "LSP header block exceeded 8 KiB without terminator"
            )
        line = line[:-2]  # strip CRLF
        if not line:
            break  # blank line ends header block
        try:
            key, _, value = line.decode("ascii").partition(":")
        except UnicodeDecodeError as e:
            raise LSPProtocolError(f"non-ASCII LSP header: {line!r}") from e
        if not key:
            raise LSPProtocolError(f"malformed LSP header line: {line!r}")
        headers[key.strip().lower()] = value.strip()

    cl = headers.get("content-length")
    if cl is None:
        raise LSPProtocolError(f"LSP message missing Content-Length: {headers!r}")
    try:
        n = int(cl)
    except ValueError as e:
        raise LSPProtocolError(f"non-integer Content-Length: {cl!r}") from e
    if n < 0 or n > _MAX_BODY_BYTES:
        raise LSPProtocolError(f"unreasonable Content-Length: {n}")

    try:
        body = await reader.readexactly(n)
    except asyncio.IncompleteReadError as e:
        raise LSPProtocolError(
            f"truncated LSP body: expected {n} bytes, got {len(e.partial)}"
        ) from e

    try:
        return json.loads(body.decode("utf-8"))
    except json.JSONDecodeError as e:
        raise LSPProtocolError(f"invalid JSON in LSP body: {e}") from e
    except UnicodeDecodeError as e:
        raise LSPProtocolError(f"non-UTF-8 LSP body: {e}") from e


def make_notification(method: str, params: Any) -> dict:
    """Build a JSON-RPC 2.0 notification envelope (no ``id``)."""
    msg: dict = {"jsonrpc": "2.0", "method": method}
    if params is not None:
        msg["params"] = params
    return msg


def make_request(req_id: int, method: str, params: Any) -> dict:
    """Build a JSON-RPC 2.0 request envelope."""
    msg: dict = {"jsonrpc": "2.0", "id": req_id, "method": method}
    if params is not None:
        msg["params"] = params
    return msg


def make_response(req_id: Any, result: Any) -> dict:
    """Build a JSON-RPC 2.0 success response envelope."""
    return {"jsonrpc": "2.0", "id": req_id, "result": result}


def make_error_response(req_id: Any, code: int, message: str, data: Any = None) -> dict:
    """Build a JSON-RPC 2.0 error response envelope."""
    err: dict = {"code": code, "message": message}
    if data is not None:
        err["data"] = data
    return {"jsonrpc": "2.0", "id": req_id, "error": err}


def classify_message(msg: dict) -> Tuple[str, Any]:
    """Return ``(kind, key)``: kind ∈ request/response/notification/invalid.

    Key is the id for request/response, the method for notifications, ``None`` for invalid.
    """
    if not isinstance(msg, dict) or msg.get("jsonrpc") != "2.0":
        return "invalid", None
    has_id = "id" in msg
    has_method = "method" in msg
    if has_id and has_method:
        return "request", msg["id"]
    if has_id and ("result" in msg or "error" in msg):
        return "response", msg["id"]
    if has_method and not has_id:
        return "notification", msg["method"]
    return "invalid", None


__all__ = [
    "ERROR_CONTENT_MODIFIED",
    "ERROR_REQUEST_CANCELLED",
    "ERROR_METHOD_NOT_FOUND",
    "LSPProtocolError",
    "LSPRequestError",
    "encode_message",
    "read_message",
    "make_request",
    "make_notification",
    "make_response",
    "make_error_response",
    "classify_message",
]