♻️ Back LSP helper with protocol library

Replace the hand-rolled JSON-RPC client (framing, pending map, reply
dispatch) with createMessageConnection and typed requests from
vscode-languageserver-protocol. A catch-all onRequest answers the server
requests the old client replied to, and CompletionList/CompletionItem
replace the ad-hoc shape guards. Public LspSession API is unchanged.

src/util/__tests__/lsp-completion.ts

@@ -4,6 +4,22 @@ import path from "node:path";
 import { setTimeout as delay } from "node:timers/promises";
 import { pathToFileURL } from "node:url";
 
+import {
+    type CompletionItem,
+    type CompletionList,
+    CompletionRequest,
+    ConfigurationRequest,
+    createMessageConnection,
+    DidOpenTextDocumentNotification,
+    ExitNotification,
+    InitializedNotification,
+    InitializeRequest,
+    type MessageConnection,
+    ShutdownRequest,
+    StreamMessageReader,
+    StreamMessageWriter,
+} from "vscode-languageserver-protocol/node";
+
 /**
  * Print the completion labels TypeScript's language server offers at a marker
  * inside a source file.
@@ -48,71 +64,18 @@ export interface CompletionResult {
     readonly labels: readonly string[];
 }
 
-interface CompletionItem {
-    readonly label: string;
-}
-
-interface RpcMessage {
-    readonly id?: string | number;
-    readonly method?: string;
-    readonly params?: unknown;
-    readonly result?: unknown;
-    readonly error?: unknown;
-}
-
-interface Deferred {
-    readonly resolve: (message: RpcMessage) => void;
-    readonly reject: (error: Error) => void;
-}
-
-const isRecord = (value: unknown): value is Record<string, unknown> =>
-    typeof value === "object" && value !== null;
-
-const isRpcMessage = (value: unknown): value is RpcMessage => {
-    if (!isRecord(value)) {
-        return false;
-    }
-    const { id, method } = value;
-    return (
-        (id === undefined ||
-            typeof id === "number" ||
-            typeof id === "string") &&
-        (method === undefined || typeof method === "string")
-    );
-};
-
-const isCompletionItem = (value: unknown): value is CompletionItem =>
-    isRecord(value) && typeof value["label"] === "string";
-
-const serverRequestResult = (message: RpcMessage): readonly null[] | null => {
-    if (message.method !== "workspace/configuration") {
-        return null;
-    }
-    const items = isRecord(message.params)
-        ? message.params["items"]
-        : undefined;
-    return Array.isArray(items) ? items.map(() => null) : [];
-};
-
-const completionLabels = (result: unknown): readonly string[] => {
-    if (!isRecord(result)) {
-        return [];
-    }
-    const { items } = result;
-    if (!Array.isArray(items)) {
+const completionLabels = (
+    result: CompletionItem[] | CompletionList | null,
+): readonly string[] => {
+    if (result === null) {
         return [];
     }
+    const items = Array.isArray(result) ? result : result.items;
     return items
-        .filter(isCompletionItem)
         .map((item) => item.label)
         .toSorted((left, right) => left.localeCompare(right));
 };
 
-const frame = (message: object): string => {
-    const body = JSON.stringify(message);
-    return `Content-Length: ${Buffer.byteLength(body)}\r\n\r\n${body}`;
-};
-
 const markerPosition = (source: string, marker: string): Position => {
     const index = source.indexOf(marker);
     if (index === -1) {
@@ -124,153 +87,81 @@ const markerPosition = (source: string, marker: string): Position => {
     return { line: lineIndex, character: line === undefined ? 0 : line.length };
 };
 
-/** Minimal JSON-RPC client for `tsc --lsp --stdio`. */
-class LspClient {
-    readonly #child: ChildProcessWithoutNullStreams;
-    readonly #pending = new Map<string | number, Deferred>();
-    #buffer = Buffer.alloc(0);
-    #nextId = 1;
-
-    public constructor(repoRoot: string) {
-        this.#child = spawn(
-            path.join(repoRoot, "node_modules", ".bin", "tsc"),
-            ["--lsp", "--stdio"],
-            { cwd: repoRoot },
-        );
-        this.#child.stdout.on("data", (chunk: Buffer) => {
-            this.#consume(chunk);
-        });
-        this.#child.stderr.on("data", (chunk: Buffer) => {
-            process.stderr.write(chunk);
-        });
-    }
-
-    /** Send a request and settle when the server answers it. */
-    public request(method: string, params?: unknown): Promise<RpcMessage> {
-        const id = this.#nextId;
-        this.#nextId += 1;
-        const message = frame({ jsonrpc: "2.0", id, method, params });
-        return new Promise<RpcMessage>((resolve, reject) => {
-            this.#pending.set(id, { resolve, reject });
-            this.#child.stdin.write(message);
-        });
-    }
-
-    public notify(method: string, params: unknown): void {
-        this.#child.stdin.write(frame({ jsonrpc: "2.0", method, params }));
-    }
-
-    public dispose(): void {
-        this.#child.stdin.end();
-        this.#child.kill();
-    }
-
-    #consume(chunk: Buffer): void {
-        this.#buffer = Buffer.concat([this.#buffer, chunk]);
-        for (;;) {
-            const headerEnd = this.#buffer.indexOf("\r\n\r\n");
-            if (headerEnd === -1) {
-                return;
-            }
-            const header = this.#buffer.subarray(0, headerEnd).toString();
-            const match = /Content-Length: (\d+)/i.exec(header);
-            const lengthText = match === null ? undefined : match[1];
-            if (lengthText === undefined) {
-                return;
-            }
-            const bodyEnd = headerEnd + 4 + Number(lengthText);
-            if (this.#buffer.length < bodyEnd) {
-                return;
-            }
-            const body = this.#buffer
-                .subarray(headerEnd + 4, bodyEnd)
-                .toString();
-            this.#buffer = this.#buffer.subarray(bodyEnd);
-            const parsed: unknown = JSON.parse(body);
-            if (isRpcMessage(parsed)) {
-                this.#handle(parsed);
-            }
-        }
-    }
-
-    #handle(message: RpcMessage): void {
-        // Server -> client request: answer so the server does not stall.
-        if (message.method !== undefined && message.id !== undefined) {
-            this.#child.stdin.write(
-                frame({
-                    jsonrpc: "2.0",
-                    id: message.id,
-                    result: serverRequestResult(message),
-                }),
-            );
-            return;
-        }
-        // Response to one of our requests.
-        const { id } = message;
-        if (id === undefined) {
-            return;
-        }
-        const deferred = this.#pending.get(id);
-        if (deferred === undefined) {
-            return;
-        }
-        this.#pending.delete(id);
-        if (message.error === undefined) {
-            deferred.resolve(message);
-        } else {
-            deferred.reject(new Error(JSON.stringify(message.error)));
-        }
-    }
-}
-
 /**
  * One language server, initialized on first use, shared across probes. Callers
  * own the lifecycle and must `close()` it.
  */
 export class LspSession {
-    readonly #client: LspClient;
+    readonly #child: ChildProcessWithoutNullStreams;
+    readonly #connection: MessageConnection;
     readonly #repoRoot: string;
     #ready: Promise<void> | undefined;
 
     public constructor(repoRoot: string) {
         this.#repoRoot = repoRoot;
-        this.#client = new LspClient(repoRoot);
+        this.#child = spawn(
+            path.join(repoRoot, "node_modules", ".bin", "tsc"),
+            ["--lsp", "--stdio"],
+            { cwd: repoRoot },
+        );
+        this.#child.stderr.on("data", (chunk: Buffer) => {
+            process.stderr.write(chunk);
+        });
+        this.#connection = createMessageConnection(
+            new StreamMessageReader(this.#child.stdout),
+            new StreamMessageWriter(this.#child.stdin),
+        );
+        // Server -> client requests the server waits on: answer so it proceeds.
+        // `workspace/configuration` wants one reply per requested item; a
+        // catch-all covers the rest (e.g. `client/registerCapability`).
+        this.#connection.onRequest(ConfigurationRequest.type, ({ items }) =>
+            items.map(() => null),
+        );
+        this.#connection.onRequest(() => null);
+        this.#connection.listen();
     }
 
     public completionLabelsAt(
         target: CompletionTarget,
     ): Promise<CompletionResult> {
-        return this.#ensureInitialized().then(() => {
-            const { uri, position } = this.#open(target);
-            return this.#client
-                .request("textDocument/completion", {
+        return this.#ensureInitialized()
+            .then(() => this.#open(target))
+            .then(({ uri, position }) =>
+                this.#connection
+                    .sendRequest(CompletionRequest.type, {
                         textDocument: { uri },
                         position,
                         context: { triggerKind: 1 },
                     })
-                .then((response) => ({
+                    .then((result) => ({
                         position,
-                    labels: completionLabels(response.result),
-                }));
-        });
+                        labels: completionLabels(result),
+                    })),
+            );
     }
 
     /** `shutdown` + `exit`, then kill the server; safe after a failed probe. */
     public close(): Promise<void> {
-        return this.#client
-            .request("shutdown")
-            .then(() => {
-                this.#client.notify("exit", null);
-                return delay(EXIT_DELAY_MS);
-            })
+        return this.#connection
+            .sendRequest(ShutdownRequest.type)
+            .then(() =>
+                this.#connection.sendNotification(ExitNotification.type),
+            )
+            .then(() => delay(EXIT_DELAY_MS))
             .finally(() => {
-                this.#client.dispose();
+                this.#connection.dispose();
+                this.#child.kill();
             });
     }
 
     #ensureInitialized(): Promise<void> {
-        this.#ready ??= this.#client
-            .request("initialize", {
+        this.#ready ??= this.#initialize();
+        return this.#ready;
+    }
+
+    #initialize(): Promise<void> {
+        return this.#connection
+            .sendRequest(InitializeRequest.type, {
                 processId: process.pid,
                 rootUri: pathToFileURL(this.#repoRoot).href,
                 workspaceFolders: [
@@ -285,16 +176,18 @@ export class LspSession {
                     },
                 },
             })
-            .then(() => {
-                this.#client.notify("initialized", {});
-            });
-        return this.#ready;
+            .then(() =>
+                this.#connection.sendNotification(
+                    InitializedNotification.type,
+                    {},
+                ),
+            );
     }
 
-    #open(target: CompletionTarget): {
+    #open(target: CompletionTarget): Promise<{
         readonly uri: string;
         readonly position: Position;
-    } {
+    }> {
         const absolute = path.resolve(this.#repoRoot, target.file);
         const marker = target.marker ?? DEFAULT_MARKER;
         const position = markerPosition(target.source, marker);
@@ -302,15 +195,16 @@ export class LspSession {
         const uri = pathToFileURL(absolute).href;
         // The server handles `didOpen` in order before the completion request,
         // so no settle delay is needed.
-        this.#client.notify("textDocument/didOpen", {
+        return this.#connection
+            .sendNotification(DidOpenTextDocumentNotification.type, {
                 textDocument: {
                     uri,
                     languageId: "typescript",
                     version: 1,
                     text,
                 },
-        });
-        return { uri, position };
+            })
+            .then(() => ({ uri, position }));
     }
 }