pi-lsp/src/client.ts

import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process";
import * as fs from "node:fs";
import {
  createMessageConnection,
  StreamMessageReader,
  StreamMessageWriter,
  type MessageConnection,
} from "vscode-jsonrpc/node.js";
import type {
  InitializeParams,
  PublishDiagnosticsParams,
} from "vscode-languageserver-protocol";
import * as path from "node:path";
import type { ServerConfig } from "./types.ts";
import { findRoot, pathToUri, uriToPath } from "./root.ts";

// Is On PATH - Returns true if `cmd` resolves to an executable via the
// current PATH. Absolute/relative paths are checked directly.
function isOnPath(cmd: string): boolean {
  const isExec = (p: string) => {
    try {
      fs.accessSync(p, fs.constants.X_OK);
      return true;
    } catch {
      return false;
    }
  };
  if (cmd.includes(path.sep)) return isExec(cmd);
  const exts =
    process.platform === "win32"
      ? (process.env.PATHEXT ?? ".EXE;.CMD;.BAT").split(";")
      : [""];
  for (const dir of (process.env.PATH ?? "").split(path.delimiter)) {
    if (!dir) continue;
    for (const ext of exts) {
      if (isExec(path.join(dir, cmd + ext))) return true;
    }
  }
  return false;
}

// LspClient - Thin wrapper that spawns a language server, performs the
// initialize handshake, auto-opens a file, and exposes sendRequest so the
// command dispatcher can forward raw LSP method calls.
//
// Method Names As Strings - We use plain method strings rather than the
// typed ProtocolRequestType constants to sidestep a version skew between
// vscode-jsonrpc and the copy re-exported by vscode-languageserver-protocol.
//
// Lifetime - One client per CLI invocation. On dispose() we do a best-effort
// shutdown/exit. A future daemon will reuse a client per (server.id, rootUri).
export class LspClient {
  private proc!: ChildProcessWithoutNullStreams;
  private conn!: MessageConnection;
  private diagnostics = new Map<string, PublishDiagnosticsParams>();
  private diagListeners = new Set<(p: PublishDiagnosticsParams) => void>();
  // Progress Tracking - Servers like gopls aren't ready to serve requests
  // until they finish their initial workspace load. They announce this via
  // `$/progress` end notifications. We track outstanding tokens so callers
  // can await readiness.
  private progressTokens = new Set<string | number>();
  private progressListeners = new Set<() => void>();

  constructor(private readonly server: ServerConfig) {}

  // Start - Spawns the server process and wires up JSON-RPC.
  async start(rootDir: string): Promise<void> {
    // Verify Binary On PATH - Fail fast with a clear message instead of
    // letting spawn ENOENT surface as a generic error. It's the user's
    // responsibility to have the server installed & on PATH.
    if (!isOnPath(this.server.command)) {
      throw new Error(
        `LSP server binary "${this.server.command}" not found on PATH. ` +
          `Install it and ensure it's on your PATH (required by server "${this.server.id}").`,
      );
    }
    this.proc = spawn(this.server.command, this.server.args, {
      stdio: ["pipe", "pipe", "pipe"],
      cwd: rootDir,
    });
    this.proc.on("error", (err) => {
      process.stderr.write(
        `[lsp:${this.server.id}] spawn error: ${err.message}\n`,
      );
    });
    // Drain Server Stderr - Many servers log verbosely; only forward when
    // LSP_DEBUG is set so normal output stays clean.
    this.proc.stderr.on("data", (chunk) => {
      if (process.env.LSP_DEBUG) process.stderr.write(chunk);
    });

    this.conn = createMessageConnection(
      new StreamMessageReader(this.proc.stdout),
      new StreamMessageWriter(this.proc.stdin),
    );

    // Capture Diagnostics - Stored by URI, also fanned out to listeners so
    // the diagnostics command can await the first publish.
    this.conn.onNotification(
      "textDocument/publishDiagnostics",
      (p: PublishDiagnosticsParams) => {
        this.diagnostics.set(p.uri, p);
        for (const l of this.diagListeners) l(p);
      },
    );

    // Track Progress Tokens - Count begin/end so we know when the server
    // has finished all startup work. We accept every createWorkDoneProgress
    // request so servers (gopls) will actually send progress.
    this.conn.onRequest(
      "window/workDoneProgress/create",
      (p: { token: string | number }) => {
        this.progressTokens.add(p.token);
        return null;
      },
    );
    this.conn.onNotification(
      "$/progress",
      (p: { token: string | number; value: { kind: string } }) => {
        if (p.value?.kind === "begin") this.progressTokens.add(p.token);
        else if (p.value?.kind === "end") {
          this.progressTokens.delete(p.token);
          for (const l of this.progressListeners) l();
        }
      },
    );
    // Accept Common Server Requests - Return empty/null so servers don't
    // stall. Good enough for a CLI; a real client would answer properly.
    this.conn.onRequest("workspace/configuration", () => []);
    this.conn.onRequest("client/registerCapability", () => null);
    this.conn.onRequest("client/unregisterCapability", () => null);

    this.conn.listen();

    const rootUri = pathToUri(rootDir);
    const params: InitializeParams = {
      processId: process.pid,
      rootUri,
      workspaceFolders: [{ uri: rootUri, name: rootDir }],
      capabilities: {
        textDocument: {
          hover: { contentFormat: ["markdown", "plaintext"] },
          definition: { linkSupport: false },
          references: {},
          completion: { completionItem: { snippetSupport: false } },
          documentSymbol: { hierarchicalDocumentSymbolSupport: true },
          publishDiagnostics: {},
          synchronization: { didSave: true },
        },
        workspace: { workspaceFolders: true, configuration: true },
      },
    };
    await this.conn.sendRequest("initialize", {
      ...params,
      capabilities: {
        ...params.capabilities,
        window: { workDoneProgress: true },
      },
    });
    this.conn.sendNotification("initialized", {});
  }

  // Wait For Ready - Resolves when there are no outstanding progress
  // tokens, or after `timeoutMs`. For servers that never send progress,
  // this effectively just waits for `graceMs` of quiet.
  async waitForReady(timeoutMs = 15000, graceMs = 300): Promise<void> {
    const deadline = Date.now() + timeoutMs;
    // Initial Grace - Give the server a moment to announce begin tokens.
    await new Promise((r) => setTimeout(r, graceMs));
    while (this.progressTokens.size > 0 && Date.now() < deadline) {
      await new Promise<void>((resolve) => {
        const onDone = () => {
          clearTimeout(timer);
          this.progressListeners.delete(onDone);
          resolve();
        };
        const timer = setTimeout(onDone, Math.min(500, deadline - Date.now()));
        this.progressListeners.add(onDone);
      });
    }
  }

  // Open Document - Reads the file from disk and sends didOpen. Most
  // servers require this before they'll answer hover/definition/etc.
  openDocument(filePath: string): string {
    const uri = pathToUri(filePath);
    const text = fs.readFileSync(filePath, "utf8");
    this.conn.sendNotification("textDocument/didOpen", {
      textDocument: {
        uri,
        languageId: this.server.languageId ?? this.server.match[0],
        version: 1,
        text,
      },
    });
    return uri;
  }

  // Send Raw LSP Request - Passthrough used by the command dispatcher.
  sendRequest<R = unknown>(method: string, params: unknown): Promise<R> {
    return this.conn.sendRequest(method, params) as Promise<R>;
  }

  // Wait For Diagnostics - Resolves on the first publish for `uri` or
  // after `timeoutMs`. Returns whatever we have for that URI.
  async waitForDiagnostics(
    uri: string,
    timeoutMs: number,
  ): Promise<PublishDiagnosticsParams> {
    if (this.diagnostics.has(uri)) return this.diagnostics.get(uri)!;
    return new Promise((resolve) => {
      const timer = setTimeout(() => {
        this.diagListeners.delete(listener);
        resolve(this.diagnostics.get(uri) ?? { uri, diagnostics: [] });
      }, timeoutMs);
      const listener = (p: PublishDiagnosticsParams) => {
        if (p.uri !== uri) return;
        clearTimeout(timer);
        this.diagListeners.delete(listener);
        resolve(p);
      };
      this.diagListeners.add(listener);
    });
  }

  // Dispose - Best-effort shutdown; kills the process if it doesn't exit.
  async dispose(): Promise<void> {
    if (this.conn) {
      try {
        await this.conn.sendRequest("shutdown", undefined);
        this.conn.sendNotification("exit");
      } catch {
        // Ignore - we're tearing down anyway.
      }
      this.conn.dispose();
    }
    if (this.proc && !this.proc.killed) {
      this.proc.kill();
    }
  }
}

// Convenience - Build a client, find root, start, and open the file.
export async function startClientForFile(
  server: ServerConfig,
  filePath: string,
): Promise<{ client: LspClient; uri: string; rootDir: string }> {
  const rootDir = findRoot(filePath, server.rootMarkers);
  const client = new LspClient(server);
  await client.start(rootDir);
  const uri = client.openDocument(filePath);
  // Wait For Workspace Load - gopls & friends reject requests with errors
  // like "no views" until their initial load completes.
  await client.waitForReady();
  return { client, uri, rootDir };
}

export { uriToPath };