cleveragents-core/benchmarks/lsp_stub_bench.py

"""ASV benchmarks for LSP server stub startup latency.

Measures the performance of:
- LspServer construction overhead
- Initialize request handling latency
- Full handshake (initialize + shutdown + exit) latency
- MethodNotFound error response latency
"""

from __future__ import annotations

import importlib
import io
import json
import sys
from pathlib import Path

# Ensure the local *source* tree is importable even when ASV has an
# older build of the package installed.
_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

# Force-reload so ASV picks up the source tree version.  Reload the
# specific submodule that contains the benchmarked code, not just the
# top-level package.
import cleveragents  # noqa: E402
import cleveragents.lsp.server  # noqa: E402

importlib.reload(cleveragents)
importlib.reload(cleveragents.lsp.server)

from cleveragents.lsp.server import LspServer  # noqa: E402


def _encode_message(msg: dict) -> bytes:
    """Encode a JSON-RPC message with Content-Length framing."""
    payload = json.dumps(msg).encode("utf-8")
    header = f"Content-Length: {len(payload)}\r\n\r\n".encode()
    return header + payload


def _make_streams(*messages: dict) -> tuple[io.BytesIO, io.BytesIO]:
    """Build input/output streams from a sequence of JSON-RPC messages."""
    raw = b"".join(_encode_message(msg) for msg in messages)
    return io.BytesIO(raw), io.BytesIO()


class LspStubStartupSuite:
    """Benchmark LSP server stub startup and request handling."""

    timeout = 60

    def setup(self) -> None:
        """Pre-build common messages and reusable streams for benchmarks."""
        self.init_msg = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {"capabilities": {}},
        }
        self.shutdown_msg = {
            "jsonrpc": "2.0",
            "id": 2,
            "method": "shutdown",
        }
        self.exit_msg = {
            "jsonrpc": "2.0",
            "method": "exit",
        }
        self.completion_msg = {
            "jsonrpc": "2.0",
            "id": 3,
            "method": "textDocument/completion",
            "params": {},
        }

        # Pre-build the 100-message payload so that list construction
        # overhead is excluded from the benchmark measurement.
        batch_messages: list[dict] = [self.init_msg]
        for i in range(100):
            batch_messages.append(
                {
                    "jsonrpc": "2.0",
                    "id": i + 10,
                    "method": "textDocument/completion",
                    "params": {},
                }
            )
        batch_messages.append(self.exit_msg)
        self._batch_raw = b"".join(_encode_message(msg) for msg in batch_messages)

    def time_server_construction(self) -> None:
        """Measure LspServer construction overhead."""
        inp, out = _make_streams(self.exit_msg)
        LspServer(input_stream=inp, output_stream=out)

    def time_initialize_request(self) -> None:
        """Measure initialize request handling latency."""
        inp, out = _make_streams(self.init_msg, self.exit_msg)
        server = LspServer(input_stream=inp, output_stream=out)
        server.run()

    def time_full_handshake(self) -> None:
        """Measure full lifecycle: initialize + shutdown + exit."""
        inp, out = _make_streams(self.init_msg, self.shutdown_msg, self.exit_msg)
        server = LspServer(input_stream=inp, output_stream=out)
        server.run()

    def time_method_not_found(self) -> None:
        """Measure MethodNotFound error response latency (after initialize)."""
        inp, out = _make_streams(self.init_msg, self.completion_msg, self.exit_msg)
        server = LspServer(input_stream=inp, output_stream=out)
        server.run()

    def time_100_method_not_found(self) -> None:
        """Measure 100 MethodNotFound responses (after initialize)."""
        inp = io.BytesIO(self._batch_raw)
        out = io.BytesIO()
        server = LspServer(input_stream=inp, output_stream=out)
        server.run()